code/08-neural-nets/overfitting.py

"""过拟合：训练集上越来越好，没见过的数据上却不见长进。以及几种对付它的办法。

故意制造过拟合：只用 100 张图训练一个很大的网络。
    python overfitting.py
全部在 CPU 上运行，十几秒。固定了随机种子。
"""
import torch
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from torch import nn

digits = load_digits()
X = torch.tensor(digits.data / 16.0, dtype=torch.float32)
y = torch.tensor(digits.target, dtype=torch.long)
# 先分出验证集（调参时看）和测试集（最后才看），剩下的是可以拿来训练的图
X_pool, X_rest, y_pool, y_rest = train_test_split(X, y, test_size=0.5, random_state=0)
X_val, X_test, y_val, y_test = train_test_split(X_rest, y_rest, test_size=0.5, random_state=0)
X_small, y_small = X_pool[:100], y_pool[:100]  # 故意只用 100 张
print(f"可用于训练的图 {len(X_pool)} 张（先只用其中 {len(X_small)} 张），验证集 {len(X_val)} 张，测试集 {len(X_test)} 张")


def make_model(dropout=0.0):
    return nn.Sequential(
        nn.Linear(64, 512), nn.ReLU(), nn.Dropout(dropout),
        nn.Linear(512, 512), nn.ReLU(), nn.Dropout(dropout),
        nn.Linear(512, 10),
    )


loss_fn = nn.CrossEntropyLoss()


def evaluate(model, X, y):
    model.eval()  # 评估模式：dropout 在这个模式下不起作用
    with torch.no_grad():
        out = model(X)
        return loss_fn(out, y).item(), (out.argmax(1) == y).float().mean().item()


def train(name, X_train=X_small, y_train=y_small, dropout=0.0, weight_decay=0.0,
          epochs=300, early_stop=False, report=()):
    torch.manual_seed(0)
    model = make_model(dropout)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.001, weight_decay=weight_decay)
    best_val, best_state, best_epoch, patience = float("inf"), None, 0, 0
    print(f"\n== {name}")
    for epoch in range(1, epochs + 1):
        model.train()  # 训练模式：dropout 生效
        loss = loss_fn(model(X_train), y_train)  # 数据不多，一次全部喂进去
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        val_loss, val_acc = evaluate(model, X_val, y_val)
        if epoch in report:
            train_loss, train_acc = evaluate(model, X_train, y_train)
            print(f"  第 {epoch:3d} 轮  训练损失 {train_loss:.4f} 准确率 {train_acc:5.1%}  |  "
                  f"验证损失 {val_loss:.4f} 准确率 {val_acc:.1%}")
        if val_loss < best_val:
            best_val, best_epoch, patience = val_loss, epoch, 0
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
        else:
            patience += 1
        if early_stop and patience >= 20:  # 验证损失连续 20 轮没有创新低，就停下来
            print(f"  第 {epoch} 轮停止：验证损失从第 {best_epoch} 轮之后就没再下降")
            break
    if early_stop:
        model.load_state_dict(best_state)  # 回到验证损失最低的那一刻
    test_loss, test_acc = evaluate(model, X_test, y_test)
    print(f"  最终：测试集准确率 {test_acc:.1%}，测试损失 {test_loss:.4f}")
    return test_acc, test_loss


print(f"模型共 {sum(p.numel() for p in make_model().parameters())} 个参数")
REPORT = (1, 10, 25, 50, 100, 200, 300)
results = {
    "什么都不加": train("什么都不加", report=REPORT),
    "权重衰减 0.1": train("权重衰减 weight_decay=0.1", weight_decay=0.1, report=(50, 300)),
    "权重衰减 1.0": train("权重衰减 weight_decay=1.0", weight_decay=1.0, report=(50, 300)),
    "dropout 0.5": train("dropout 0.5", dropout=0.5, report=(50, 300)),
    "早停": train("早停", early_stop=True),
    f"{len(X_pool)} 张训练数据": train(f"什么都不加，但用 {len(X_pool)} 张图训练", X_pool, y_pool, report=REPORT),
}
print("\n测试集汇总：")
for k, (acc, loss) in results.items():
    print(f"  {k:<16} 准确率 {acc:.1%}  损失 {loss:.4f}")