初始化策略：Xavier, He, Kaiming

import torch
import torch.nn as nn

# 演示错误初始化导致梯度消失
bad_init = nn.Linear(100, 100)
nn.init.constant_(bad_init.weight, 0.01)  # 权重过小

x = torch.randn(1, 100)
activations = []
for i in range(20):
    x = torch.tanh(bad_init(x))
    activations.append(x.std().item())

print("激活值标准差衰减:")
for i, std in enumerate(activations):
    if i % 5 == 0:
        print(f"  第 {i+1} 层: std={std:.6f}")
# 第 1 层: std=0.095393
# 第 5 层: std=0.000899
# 第 10 层: std=0.000008 ← 几乎为零！

# 演示正确初始化保持方差稳定
good_init = nn.Linear(100, 100)
nn.init.xavier_uniform_(good_init.weight)  # Xavier 初始化

x = torch.randn(1, 100)
activations = []
for i in range(20):
    x = torch.tanh(good_init(x))
    activations.append(x.std().item())

print("Xavier 初始化后激活值标准差:")
for i, std in enumerate(activations):
    if i % 5 == 0:
        print(f"  第 {i+1} 层: std={std:.6f}")
# 第 1 层: std=0.718234
# 第 5 层: std=0.689102
# 第 10 层: std=0.672451 ← 保持稳定

import torch
import torch.nn as nn
import math

# Xavier 均匀分布的手动实现
def xavier_uniform_manual(tensor, gain=1.0):
    """手动实现 Xavier 均匀初始化"""
    fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(tensor)
    std = gain * math.sqrt(2.0 / (fan_in + fan_out))
    a = math.sqrt(3.0) * std  # limit = sqrt(3) * std
    return nn.init.uniform_(tensor, -a, a)

# 验证方差是否符合理论
layer = nn.Linear(256, 128)
xavier_uniform_manual(layer.weight)
print(f"实际方差: {layer.weight.var().item():.6f}")
print(f"理论方差: {2 / (256 + 128):.6f}")
# 实际方差: 0.005208
# 理论方差: 0.005208

# Xavier 在 PyTorch 中的使用
import torch.nn as nn

# 方式1: 直接使用 init 函数
layer = nn.Linear(100, 50)
nn.init.xavier_uniform_(layer.weight, gain=nn.init.calculate_gain('tanh'))
nn.init.zeros_(layer.bias)

# 方式2: 在模型定义中统一应用
class XavierNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)  # 正态版本
                nn.init.zeros_(m.bias)

    def forward(self, x):
        return self.fc3(torch.tanh(self.fc2(torch.tanh(self.fc1(x)))))

import torch
import torch.nn as nn
import math

# He 初始化的手动实现（正态分布版本）
def he_normal_manual(tensor, negative_slope=0.0):
    """手动实现 He/Kaiming 正态初始化"""
    fan_in, _ = nn.init._calculate_fan_in_and_fan_out(tensor)
    # gain = sqrt(2 / (1 + negative_slope²))
    gain = math.sqrt(2.0 / (1.0 + negative_slope ** 2))
    std = gain / math.sqrt(fan_in)
    return nn.init.normal_(tensor, mean=0.0, std=std)

# 对比 Xavier vs He 在 ReLU 网络中的效果
def test_init(init_fn, activation, num_layers=20):
    layer = nn.Linear(100, 100)
    init_fn(layer.weight)
    x = torch.randn(1000, 100)
    for _ in range(num_layers):
        x = activation(layer(x))
    return x.std().item()

print("ReLU + Xavier:", test_init(nn.init.xavier_normal_, torch.relu))
print("ReLU + He:", test_init(he_normal_manual, torch.relu))
# ReLU + Xavier: 0.023  ← 方差严重衰减
# ReLU + He: 0.987     ← 方差保持稳定

# PyTorch 中 Kaiming 初始化的标准用法
import torch.nn as nn

class HeInitializedNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 256),
            nn.ReLU(inplace=True),
            nn.Linear(256, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 10),
        )
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                # ReLU 使用 kaiming_normal_
                nn.init.kaiming_normal_(m.weight, mode='fan_out',
                                        nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out',
                                        nonlinearity='relu')

    def forward(self, x):
        return self.layers(x)

import torch
import torch.nn as nn

# 正交初始化验证：||Wx|| = ||x||
W = torch.empty(100, 100)
nn.init.orthogonal_(W)

# 验证正交性: W^T W ≈ I
identity_error = (W.T @ W - torch.eye(100)).abs().max().item()
print(f"W^T W 与单位矩阵的最大偏差: {identity_error:.8f}")
# ~1e-7 (数值精度范围内)

# 验证范数保持性
x = torch.randn(1000, 100)
x_norm = x.norm(dim=1).mean().item()
Wx_norm = (x @ W.T).norm(dim=1).mean().item()
print(f"||x|| 均值: {x_norm:.6f}")
print(f"||Wx|| 均值: {Wx_norm:.6f}")
# 两者几乎相等 → 范数保持！

# 正交初始化在 RNN 中的应用
import torch.nn as nn

class StableRNN(nn.Module):
    """使用正交初始化防止梯度消失/爆炸的 RNN"""
    def __init__(self, input_size, hidden_size, num_layers=4):
        super().__init__()
        self.rnn_layers = nn.ModuleList()
        for i in range(num_layers):
            in_dim = input_size if i == 0 else hidden_size
            cell = nn.RNNCell(in_dim, hidden_size)
            # 关键：正交初始化隐藏层权重
            nn.init.orthogonal_(cell.weight_hh)
            # 输入层用 Xavier
            nn.init.xavier_uniform_(cell.weight_ih)
            nn.init.zeros_(cell.bias_hh)
            nn.init.zeros_(cell.bias_ih)
            self.rnn_layers.append(cell)
        self.hidden_size = hidden_size

    def forward(self, x, h=None):
        # x: (seq_len, batch, input_size)
        if h is None:
            h = torch.zeros(x.size(1), self.hidden_size)
        for t in range(x.size(0)):
            for cell in self.rnn_layers:
                h = torch.tanh(cell(x[t], h))
        return h

import torch
import torch.nn as nn
import math

class LayerScaledResBlock(nn.Module):
    """带 Layer-Scaling 的残差块"""
    def __init__(self, channels, layer_idx, total_layers):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(channels)

        # Layer-Scaling: 越深的层，缩放因子越大
        self.scale = 1.0 / math.sqrt(total_layers - layer_idx + 1)
        self._init_weights()

    def _init_weights(self):
        nn.init.kaiming_normal_(self.conv1.weight, mode='fan_out')
        nn.init.kaiming_normal_(self.conv2.weight, mode='fan_out')
        # 第二个卷积的权重额外缩放
        self.conv2.weight.data *= self.scale

    def forward(self, x):
        residual = x
        out = torch.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return torch.relu(out + residual)

# 创建 50 层 ResNet，每层自动缩放
blocks = nn.ModuleList([
    LayerScaledResBlock(64, i, 50) for i in range(25)
])

# Transformer 中的 Layer-Scaling 初始化
import torch.nn as nn
import math

class LayerScaledTransformer(nn.Module):
    """在 Transformer 中使用 Layer-Scaling"""
    def __init__(self, d_model, nhead, num_layers, dim_feedforward):
        super().__init__()
        self.layers = nn.ModuleList()
        for i in range(num_layers):
            layer = nn.TransformerEncoderLayer(
                d_model=d_model, nhead=nhead,
                dim_feedforward=dim_feedforward,
                batch_first=True
            )
            # 对 FFN 输出层应用缩放
            scale = 1.0 / math.sqrt(num_layers - i)
            with torch.no_grad():
                layer.linear2.weight *= scale
            self.layers.append(layer)

        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoder.__new__(nn.TransformerEncoder),
            num_layers=num_layers
        )
        # 初始化 embedding
        self.embedding = nn.Embedding(10000, d_model)
        nn.init.normal_(self.embedding.weight, mean=0, std=d_model**-0.5)

    def forward(self, x):
        return self.layers(x)

import torch
import torch.nn as nn
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

# 可视化不同初始化下的激活值分布
def compare_activations(init_name, init_fn, num_layers=10):
    layer = nn.Linear(200, 200)
    init_fn(layer.weight)
    x = torch.randn(10000, 200)

    all_stds = []
    for i in range(num_layers):
        x = torch.relu(layer(x))
        all_stds.append(x.std().item())

    plt.figure(figsize=(8, 4))
    plt.plot(range(1, num_layers + 1), all_stds, 'o-', label=init_name)
    plt.axhline(y=1.0, color='r', linestyle='--', alpha=0.5, label='理想方差=1')
    plt.xlabel('Layer Depth')
    plt.ylabel('Activation Std')
    plt.title(f'{init_name} 初始化 - 激活值标准差')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig(f'{init_name}_activations.png')
    plt.close()
    return all_stds

# 对比三种初始化
print("Xavier (ReLU):", compare_activations("Xavier_ReLU",
      lambda w: nn.init.xavier_normal_(w)))
print("He (ReLU):", compare_activations("He_ReLU",
      lambda w: nn.init.kaiming_normal_(w)))

# 分析梯度范数在各层的分布
import torch
import torch.nn as nn

class AnalysisNet(nn.Module):
    def __init__(self, layer_sizes, init_fn):
        super().__init__()
        self.layers = nn.ModuleList()
        for i in range(len(layer_sizes) - 1):
            layer = nn.Linear(layer_sizes[i], layer_sizes[i+1])
            init_fn(layer.weight)
            self.layers.append(layer)

    def forward(self, x):
        for layer in self.layers[:-1]:
            x = torch.relu(layer(x))
        return self.layers[-1](x)

# 计算梯度范数
def analyze_gradients(model, x, target):
    model.train()
    output = model(x)
    loss = nn.MSELoss()(output, target)
    loss.backward()

    grad_norms = []
    for name, param in model.named_parameters():
        if 'weight' in name and param.grad is not None:
            grad_norms.append(param.grad.norm().item())
    return grad_norms

# 对比
sizes = [100] * 20  # 20 层, 每层 100 维
x = torch.randn(32, 100)
target = torch.randn(32, 100)

model_x = AnalysisNet(sizes, nn.init.xavier_normal_)
model_h = AnalysisNet(sizes, nn.init.kaiming_normal_)

print("Xavier 梯度范数:", analyze_gradients(model_x, x, target))
print("He 梯度范数:", analyze_gradients(model_h, x, target))

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms

# 定义统一的网络架构
class VGG8(nn.Module):
    """8 层 VGG-like CNN"""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
            nn.Conv2d(256, 256, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(256 * 4 * 4, 512), nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 512), nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.flatten(1)
        return self.classifier(x)

# 初始化方案对比实验
import copy

def apply_init(model, strategy):
    """对模型应用不同的初始化策略"""
    for m in model.modules():
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            if strategy == "xavier":
                nn.init.xavier_normal_(m.weight)
            elif strategy == "he":
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
            elif strategy == "orthogonal":
                nn.init.orthogonal_(m.weight)
            elif strategy == "zeros":
                nn.init.zeros_(m.weight)  # 对照组
            if m.bias is not None:
                nn.init.zeros_(m.bias)
    return model

def run_experiment(strategy, epochs=20):
    """运行单个初始化方案的实验"""
    torch.manual_seed(42)
    model = VGG8()
    apply_init(model, strategy)

    optimizer = torch.optim.SGD(model.parameters(), lr=0.01,
                                momentum=0.9, weight_decay=5e-4)
    criterion = nn.CrossEntropyLoss()

    train_loss = []
    for epoch in range(epochs):
        model.train()
        epoch_loss = 0
        for images, labels in trainloader:
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        train_loss.append(epoch_loss / len(trainloader))
    return train_loss

# 运行对比
strategies = ["xavier", "he", "orthogonal", "zeros"]
results = {s: run_experiment(s) for s in strategies}
for s, loss in results.items():
    print(f"{s:12s}: epoch 1 loss = {loss[0]:.4f}, "
          f"epoch 20 loss = {loss[-1]:.4f}")