Weight Initialization

import torch.nn as nn

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

# He init cho ReLU (PyTorch mặc định Kaiming uniform)
for m in model.modules():
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
        nn.init.zeros_(m.bias)  # bias luôn init = 0

# Xavier init cho Sigmoid/Tanh
# nn.init.xavier_normal_(m.weight)

# Kiểm tra phương sai ban đầu
with torch.no_grad():
    x = torch.randn(100, 784)
    for layer in model:
        x = layer(x)
        if isinstance(layer, nn.Linear):
            print(f"Var = {x.var().item():.4f}")