PyTorch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])

print(x.shape)                   # torch.Size([2, 3])
print(x.sum(dim=0))               # tensor([5., 7., 9.])
print(x.sum(dim=1))               # tensor([6., 15.])
print(x.sum(dim=1, keepdim=True))  # tensor([[6.], [15.]])

x = torch.ones(2, 3)
bias = torch.tensor([1., 2., 3.])

print(x + bias)
# tensor([[2., 3., 4.],
#         [2., 3., 4.]])

x = torch.zeros(2, 3)
y = torch.ones(2, 3)

print(torch.cat([x, y], dim=0).shape)    # [4, 3]
print(torch.cat([x, y], dim=1).shape)    # [2, 6]
print(torch.stack([x, y], dim=0).shape)  # [2, 2, 3]
print(torch.stack([x, y], dim=1).shape)  # [2, 2, 3]

x = torch.tensor([1., 2., 3.])
a = x[:2]
b = x.clone()

a[0] = 9
print(x)  # tensor([9., 2., 3.])
print(b)  # tensor([1., 2., 3.])

logits = torch.tensor([[2., 0., -1.], [0., 1., 3.]])
labels = torch.tensor([0, 2], dtype=torch.long)

loss = nn.CrossEntropyLoss()(logits, labels)
pred = logits.argmax(dim=1)
print(pred)  # tensor([0, 2])

x = torch.tensor(2., requires_grad=True)

(x * 3).backward()
print(x.grad)  # tensor(3.)

(x * 4).backward()  # 新建一次计算图
print(x.grad)      # tensor(7.)

x.grad = None
(x * 4).backward()
print(x.grad)  # tensor(4.)

torch.save({
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}, "checkpoint.pth")

checkpoint = torch.load(
    "checkpoint.pth", map_location=device, weights_only=True
)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
start_epoch = checkpoint["epoch"] + 1

from pathlib import Path

import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms


class MNISTModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        return self.net(x)


def train_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total_loss = 0.0
    total = 0

    for images, labels in loader:
        images = images.to(device)
        labels = labels.to(device)
        optimizer.zero_grad(set_to_none=True)

        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * labels.size(0)
        total += labels.size(0)

    return total_loss / total


@torch.inference_mode()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images = images.to(device)
        labels = labels.to(device)
        logits = model(images)

        total_loss += criterion(logits, labels).item() * labels.size(0)
        correct += (logits.argmax(dim=1) == labels).sum().item()
        total += labels.size(0)

    return total_loss / total, correct / total


def main():
    torch.manual_seed(42)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    checkpoint_path = Path("mnist_best.pth")

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),
    ])
    # ToTensor 将像素缩放到 [0, 1]，再用 (x - 0.5) / 0.5 映射到 [-1, 1]。
    # 这里的 0.5 是选定的缩放参数，不是 MNIST 的实际均值和标准差。
    full_train = datasets.MNIST(
        "./data", train=True, download=True, transform=transform
    )
    test_set = datasets.MNIST(
        "./data", train=False, download=True, transform=transform
    )
    train_set, val_set = random_split(
        full_train, [54000, 6000],
        generator=torch.Generator().manual_seed(42),
    )

    train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
    val_loader = DataLoader(val_set, batch_size=256, shuffle=False)
    test_loader = DataLoader(test_set, batch_size=256, shuffle=False)

    model = MNISTModel().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    best_loss = float("inf")

    for epoch in range(5):
        train_loss = train_epoch(
            model, train_loader, criterion, optimizer, device
        )
        val_loss, val_acc = evaluate(model, val_loader, criterion, device)

        if val_loss < best_loss:
            best_loss = val_loss
            torch.save(model.state_dict(), checkpoint_path)

        print(
            f"Epoch {epoch + 1}: train={train_loss:.4f}, "
            f"val={val_loss:.4f}, acc={val_acc:.2%}"
        )

    # 重新构造模型，读取验证集表现最好的权重。
    best_model = MNISTModel().to(device)
    state = torch.load(
        checkpoint_path, map_location=device, weights_only=True
    )
    best_model.load_state_dict(state)
    test_loss, test_acc = evaluate(
        best_model, test_loader, criterion, device
    )
    print(f"Test: loss={test_loss:.4f}, acc={test_acc:.2%}")

    # 单张图片需要补上批次轴：[1, 28, 28] → [1, 1, 28, 28]。
    image, label = test_set[0]
    best_model.eval()
    with torch.inference_mode():
        logits = best_model(image.unsqueeze(0).to(device))
        prediction = logits.argmax(dim=1).item()
    print(f"Prediction: {prediction}, label: {label}")


if __name__ == "__main__":
    main()