训练曲线、学习率与优化器：判断模型为什么学不好

import copy

import numpy as np
import torch
from torch import nn

torch.manual_seed(42)
rng = np.random.default_rng(42)
X = rng.uniform(-2.5, 2.5, size=(500, 2)).astype("float32")
y = ((X[:, 0] ** 2 + 0.7 * X[:, 1]) > 1.4).astype("int64")

indices = rng.permutation(len(X))
train_index, valid_index = indices[:400], indices[400:]
X_train, y_train = torch.from_numpy(X[train_index]), torch.from_numpy(y[train_index])
X_valid, y_valid = torch.from_numpy(X[valid_index]), torch.from_numpy(y[valid_index])

template = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
initial_state = copy.deepcopy(template.state_dict())


def run(learning_rate: float) -> tuple[list[float], list[float]]:
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
    model.load_state_dict(initial_state)
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
    loss_fn = nn.CrossEntropyLoss()
    train_history, valid_history = [], []

    for _ in range(80):
        model.train()
        logits = model(X_train)
        loss = loss_fn(logits, y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_history.append(float(loss.item()))

        model.eval()
        with torch.inference_mode():
            valid_loss = loss_fn(model(X_valid), y_valid)
        valid_history.append(float(valid_loss.item()))
    return train_history, valid_history


for rate in (0.001, 0.05, 1.0):
    train_curve, valid_curve = run(rate)
    print(
        rate,
        "train", round(train_curve[-1], 4),
        "valid", round(valid_curve[-1], 4),
        "best valid", round(min(valid_curve), 4),
    )

import matplotlib.pyplot as plt

train_curve, valid_curve = run(0.05)
figure, axis = plt.subplots(figsize=(6, 4))
axis.plot(train_curve, label="train")
axis.plot(valid_curve, label="valid")
axis.set(xlabel="Epoch", ylabel="Cross-entropy", title="Training diagnosis")
axis.legend()
figure.tight_layout()
figure.savefig("training_curve.png", dpi=150)
plt.close(figure)