PyTorch 网络构建：Module、网络层与损失函数配对

from torch import nn


class Classifier(nn.Module):
    def __init__(self, input_size: int, hidden_size: int, classes: int):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, classes),
        )

    def forward(self, features):
        return self.network(features)

model = Classifier(4, 12, 3)
trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
print(model)
print("trainable parameters:", trainable)

optimizer.zero_grad()
logits = model(features)
loss = loss_fn(logits, labels)
loss.backward()
optimizer.step()

import torch
from torch import nn

torch.manual_seed(42)
features = torch.tensor([
    [0.2, 0.8, 0.1, 0.4],
    [0.9, 0.1, 0.6, 0.2],
    [0.4, 0.7, 0.8, 0.5],
    [0.1, 0.2, 0.2, 0.9],
], dtype=torch.float32)
labels = torch.tensor([0, 1, 2, 0], dtype=torch.long)


class Classifier(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(4, 12),
            nn.ReLU(),
            nn.Linear(12, 3),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.network(inputs)


model = Classifier()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)

model.train()
optimizer.zero_grad()
logits = model(features)
loss = loss_fn(logits, labels)
loss.backward()
optimizer.step()

print("logits shape:", logits.shape)
print("label shape:", labels.shape)
print("loss:", round(loss.item(), 4))

model.eval()
with torch.inference_mode():
    probabilities = model(features).softmax(dim=1)
    predictions = probabilities.argmax(dim=1)
print("probability sums:", probabilities.sum(dim=1))
print("predictions:", predictions.tolist())

first_weight = next(model.parameters())
print(first_weight.shape)
print("gradient exists:", first_weight.grad is not None)
print("gradient norm:", round(first_weight.grad.norm().item(), 6))