S2-03 · 数据增强：不花钱的提分手段

# s2_augmentation.py —— 500 样本上实测旋转增强
import torch, torch.nn as nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset

torch.manual_seed(42)
plain = transforms.ToTensor()
aug   = transforms.Compose([transforms.RandomRotation(10), transforms.ToTensor()])

train_raw = datasets.MNIST("data", train=True, download=False, transform=plain)
test_ds   = datasets.MNIST("data", train=False, download=False, transform=plain)
test_dl   = DataLoader(test_ds, batch_size=256)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))

def run(train_transform, tag):
    torch.manual_seed(42)                      # 每组实验重置种子，公平对比
    global model
    model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
    train_ds = datasets.MNIST("data", train=True, download=False, transform=train_transform)
    dl = DataLoader(Subset(train_ds, range(500)), batch_size=64, shuffle=True)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(6):
        for x, y in dl:
            opt.zero_grad(); loss_fn(model(x), y).backward(); opt.step()
    @torch.no_grad()
    def acc():
        n = ok = 0
        for x, y in test_dl:
            ok += (model(x).argmax(1) == y).sum().item(); n += len(y)
        return ok / n
    print(f"{tag}: 测试集准确率 = {acc():.4f}")

run(plain, "无增强(对照)")
run(aug,   "旋转±10°增强  ")