PyTorch: train, eval, Dropout và BatchNorm

from pathlib import Path
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
OUTPUT = Path(__file__).resolve().parent / "results"
OUTPUT.mkdir(exist_ok=True)
plt.rcParams.update({"font.size": 12, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.facecolor": "#f3f7f6",
                     "axes.facecolor": "white", "axes.titleweight": "bold"})
FIGURE = OUTPUT / "pytorch-train-eval.png"
import torch
from torch import nn
import torch.nn.functional as F
torch.manual_seed(7)
torch.set_num_threads(1)

x = torch.randn(16, 6) * 2 + 3
model = nn.Sequential(nn.BatchNorm1d(6), nn.Dropout(p=0.5), nn.Linear(6, 2))
model.train()
before = model[0].running_mean.clone()
train_a = model(x)
train_b = model(x)
assert not torch.allclose(train_a, train_b)
assert not torch.equal(before, model[0].running_mean)
print("Train outputs differ:", not torch.allclose(train_a, train_b))

model.eval()
running = model[0].running_mean.clone()
eval_with_grad = model(x)
assert eval_with_grad.requires_grad
with torch.inference_mode():
    eval_a = model(x)
    eval_b = model(x)
assert torch.allclose(eval_a, eval_b)
assert not eval_a.requires_grad
assert torch.equal(running, model[0].running_mean)
model.train()
assert model.training

fig, axes = plt.subplots(1, 2, figsize=(12, 8), layout="constrained")
axes[0].plot(train_a.detach()[:, 0], "o-", label="Train pass A")
axes[0].plot(train_b.detach()[:, 0], "s--", label="Train pass B")
axes[1].plot(eval_a[:, 0], "o-", label="Eval pass A")
axes[1].plot(eval_b[:, 0], "x--", label="Eval pass B")
for ax, title in zip(axes, ["Stochastic training mode", "Deterministic evaluation mode"]):
    ax.set(title=title, xlabel="Sample", ylabel="First logit")
    ax.legend()
fig.savefig(FIGURE, dpi=100)
plt.close(fig)
print("Saved:", FIGURE.name)

.\.venv\Scripts\python.exe pytorch_train_eval.py