Weight Initialization › Empirical Comparison of Initialization Methods

import torch
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset


def create_network_with_init(
    hidden_sizes, activation="relu", init_method="he_normal"
):
    """Create a sequential network with the specified initialization method."""
    layers = []
    for i in range(len(hidden_sizes) - 1):
        linear = nn.Linear(hidden_sizes[i], hidden_sizes[i + 1])

        if init_method == "xavier_uniform":
            nn.init.xavier_uniform_(linear.weight)
        elif init_method == "xavier_normal":
            nn.init.xavier_normal_(linear.weight)
        elif init_method == "he_normal":
            nn.init.kaiming_normal_(linear.weight, nonlinearity="relu")
        elif init_method == "he_uniform":
            nn.init.kaiming_uniform_(linear.weight, nonlinearity="relu")
        elif init_method == "orthogonal":
            nn.init.orthogonal_(linear.weight, gain=np.sqrt(2))
        elif init_method == "zeros":
            nn.init.zeros_(linear.weight)
        elif init_method == "bert":
            nn.init.normal_(linear.weight, mean=0.0, std=0.02)

        nn.init.zeros_(linear.bias)
        layers.append(linear)

        if i < len(hidden_sizes) - 2:
            if activation == "relu":
                layers.append(nn.ReLU())
            elif activation == "tanh":
                layers.append(nn.Tanh())

    return nn.Sequential(*layers)


def train_model(model, train_loader, epochs=60, lr=0.01):
    """Train model and return loss history."""
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr)
    losses = []
    for epoch in range(epochs):
        epoch_loss = 0
        for X, y in train_loader:
            optimizer.zero_grad()
            outputs = model(X)
            loss = criterion(outputs, y)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        losses.append(epoch_loss / len(train_loader))
    return losses


# Create a structured 5-class classification dataset
torch.manual_seed(42)
np.random.seed(42)

n_samples = 2000
n_features = 100
n_classes = 5

# Gaussian clusters with realistic overlap
centers = np.random.randn(n_classes, n_features) * 3
X_np = np.vstack(
    [
        centers[c] + np.random.randn(n_samples // n_classes, n_features) * 1.8
        for c in range(n_classes)
    ]
)
y_np = np.repeat(np.arange(n_classes), n_samples // n_classes)

X = torch.FloatTensor(X_np)
y = torch.LongTensor(y_np)
dataset = TensorDataset(X, y)
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)

hidden_sizes = [100, 128, 128, 128, n_classes]