训练稳定性：初始化、梯度裁剪、BatchNorm 与 Dropout

for name, parameter in model.named_parameters():
    if parameter.grad is not None:
        print(name, parameter.grad.norm().item())

loss.backward()
total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

import torch
from torch import nn

torch.manual_seed(7)


class StableNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(20, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Dropout(p=0.3),
            nn.Linear(32, 2),
        )
        nn.init.kaiming_normal_(self.layers[0].weight, nonlinearity="relu")
        nn.init.zeros_(self.layers[0].bias)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.layers(inputs)


model = StableNet()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=1e-3)
loss_fn = nn.CrossEntropyLoss()
features = torch.randn(64, 20) * 4
labels = torch.randint(0, 2, (64,))

model.train()
optimizer.zero_grad()
loss = loss_fn(model(features), labels)
loss.backward()
gradient_norm = nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
print("loss:", round(loss.item(), 4))
print("norm before clipping:", round(float(gradient_norm), 4))

sample = features[:4]
model.train()
train_output_1 = model(sample)
train_output_2 = model(sample)
model.eval()
with torch.inference_mode():
    eval_output_1 = model(sample)
    eval_output_2 = model(sample)
print("train outputs equal:", torch.allclose(train_output_1, train_output_2))
print("eval outputs equal:", torch.allclose(eval_output_1, eval_output_2))