【Bug已解决】How to decide which mode to use for ‘kaiming_normal‘ initialization 解决方案

【Bug已解决】How to decide which mode to use for 'kaiming_normal' initialization 解决方案

问题描述

在 PyTorch 中使用kaiming_normal_（也称为 He 初始化）进行权重初始化时，开发者经常对mode参数的选择感到困惑。mode参数有两个选项：'fan_in'和'fan_out'，选择错误的模式可能导致训练不稳定、收敛缓慢甚至梯度爆炸/消失。

import torch import torch.nn as nn # kaiming_normal_ 的签名 # torch.nn.init.kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu') # 常见的困惑： # 1. mode='fan_in' 和 mode='fan_out' 有什么区别？ # 2. 什么时候该用 fan_in，什么时候该用 fan_out？ # 3. nonlinearity 参数如何选择？ # 4. a 参数代表什么？

mode参数决定了初始化时使用的归一化因子，直接影响权重的方差大小，进而影响前向传播和反向传播中信号的方差。

错误复现

问题 1：使用错误的 mode 导致梯度爆炸

import torch import torch.nn as nn # 深层网络，使用错误的初始化 class DeepNetwork(nn.Module): def __init__(self, depth=20, width=256, init_mode='fan_out'): super().__init__() layers = [] layers.append(nn.Linear(784, width)) for _ in range(depth - 2): layers.append(nn.Linear(width, width)) layers.append(nn.Linear(width, 10)) self.layers = nn.ModuleList(layers) # 使用指定 mode 初始化 for layer in self.layers: nn.init.kaiming_normal_(layer.weight, mode=init_mode, nonlinearity='relu') nn.init.zeros_(layer.bias) def forward(self, x): for i, layer in enumerate(self.layers): x = layer(x) if i < len(self.layers) - 1: x = torch.relu(x) return x # 测试不同 mode 的效果 x = torch.randn(32, 784) # fan_in 模式（推荐用于前向传播） model_in = DeepNetwork(depth=20, init_mode='fan_in') out_in = model_in(x) print(f"fan_in output std: {out_in.std().item():.4f}") # fan_out 模式（推荐用于反向传播） model_out = DeepNetwork(depth=20, init_mode='fan_out') out_out = model_out(x) print(f"fan_out output std: {out_out.std().item():.4f}") # 对比梯度 loss_in = model_in(x).sum() loss_in.backward() grad_in = model_in.layers[0].weight.grad.std().item() loss_out = model_out(x).sum() loss_out.backward() grad_out = model_out.layers[0].weight.grad.std().item() print(f"fan_in gradient std: {grad_in:.6f}") print(f"fan_out gradient std: {grad_out:.6f}")

问题 2：nonlinearity 不匹配

import torch import torch.nn as nn # 错误：使用 ReLU 激活函数但 nonlinearity 设为 'leaky_relu' layer = nn.Linear(100, 200) nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='leaky_relu') # 但实际使用 ReLU，导致初始化方差不正确 # 正确：nonlinearity 应与实际激活函数匹配 nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='relu')

根因分析

1. Kaiming 初始化的数学原理

Kaiming 初始化（He et al., 2015）的目标是：在深层网络中保持每一层输出的方差与输入方差一致，避免信号在前向传播或反向传播中指数级放大或衰减。

对于权重矩阵W，Kaiming 初始化从以下分布中采样：

W ~ N(0, Var(W))

其中Var(W)的计算取决于mode：

mode='fan_in'（保持前向传播方差）：

Var(W) = 2 / fan_in

其中fan_in = weight.size(1) * receptive_field_size（输入神经元数量）。

mode='fan_out'（保持反向传播方差）：

Var(W) = 2 / fan_out

其中fan_out = weight.size(0) * receptive_field_size（输出神经元数量）。

2. fan_in vs fan_out 的选择原则

• fan_in：保持前向传播中信号方差的稳定。适用于大多数场景，特别是关注前向传播性能时。PyTorch 默认使用fan_in。

• fan_out：保持反向传播中梯度方差的稳定。适用于需要特别关注梯度流动的场景，如在非常深的网络中防止梯度消失。

对于全连接层nn.Linear(in_features, out_features)：

• fan_in = in_features
• fan_out = out_features

当in_features == out_features时，两种模式没有区别。当它们不同时，选择会影响权重的方差。

3. nonlinearity 参数

nonlinearity参数决定了增益（gain）的计算：

a参数是 LeakyReLU 的负斜率，默认为 0（即标准 ReLU）。

4. PyTorch 默认初始化的问题

PyTorch 的nn.Linear和nn.Conv2d默认使用 Kaiming uniform 初始化（kaiming_uniform_witha=sqrt(5)），这在某些情况下不是最优的。手动指定kaiming_normal_可以获得更好的训练效果。

解决方案

方案一：根据场景选择 mode

import torch import torch.nn as nn def init_weights(model: nn.Module, mode: str = 'fan_in'): """ 统一初始化模型权重 mode='fan_in': 保持前向传播方差（推荐，默认） mode='fan_out': 保持反向传播梯度方差 """ for m in model.modules(): if isinstance(m, nn.Linear): # 全连接层: nonlinearity='relu' 适用于 ReLU 激活 nn.init.kaiming_normal_(m.weight, mode=mode, nonlinearity='relu') if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.Conv2d): # 卷积层: nonlinearity='relu' 适用于 ReLU 激活 nn.init.kaiming_normal_(m.weight, mode=mode, nonlinearity='relu') if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): # BatchNorm: weight=1, bias=0 nn.init.ones_(m.weight) nn.init.zeros_(m.bias) # 使用示例 model = nn.Sequential( nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10) ) # 使用 fan_in 初始化（推荐） init_weights(model, mode='fan_in')

方案二：根据激活函数选择 nonlinearity

import torch import torch.nn as nn import torch.nn.functional as F class FlexibleInitNet(nn.Module): """支持不同激活函数的网络，自动匹配初始化""" def __init__(self, activation='relu'): super().__init__() self.activation = activation 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): """根据激活函数选择正确的 nonlinearity 参数""" nonlinearity_map = { 'relu': 'relu', 'leaky_relu': 'leaky_relu', 'elu': 'relu', # ELU 近似用 relu 'gelu': 'relu', # GELU 近似用 relu 'selu': 'linear', # SELU 有自己的初始化 'tanh': 'tanh', 'sigmoid': 'sigmoid', } nonlinearity = nonlinearity_map.get(self.activation, 'relu') a = 0.01 if self.activation == 'leaky_relu' else 0 for layer in [self.fc1, self.fc2, self.fc3]: if self.activation == 'selu': # SELU 使用特殊的初始化 nn.init.kaiming_normal_(layer.weight, mode='fan_in', nonlinearity='linear') # SELU 推荐使用 lecun_normal fan_in = layer.weight.size(1) std = (1.0 / fan_in) ** 0.5 nn.init.normal_(layer.weight, mean=0, std=std) else: nn.init.kaiming_normal_(layer.weight, a=a, mode='fan_in', nonlinearity=nonlinearity) nn.init.zeros_(layer.bias) def forward(self, x): x = self._activate(self.fc1(x)) x = self._activate(self.fc2(x)) return self.fc3(x) def _activate(self, x): if self.activation == 'relu': return F.relu(x) elif self.activation == 'leaky_relu': return F.leaky_relu(x, 0.01) elif self.activation == 'elu': return F.elu(x) elif self.activation == 'gelu': return F.gelu(x) elif self.activation == 'selu': return F.selu(x) elif self.activation == 'tanh': return torch.tanh(x) elif self.activation == 'sigmoid': return torch.sigmoid(x) return x # 使用示例 for act in ['relu', 'leaky_relu', 'tanh', 'sigmoid']: model = FlexibleInitNet(activation=act) x = torch.randn(32, 784) out = model(x) print(f"{act}: output std={out.std().item():.4f}")

方案三：对比 fan_in 和 fan_out 的效果

import torch import torch.nn as nn def compare_init_modes(depth=30, width=256, input_dim=784, output_dim=10): """对比 fan_in 和 fan_out 在深层网络中的效果""" def build_model(mode): layers = [nn.Linear(input_dim, width)] for _ in range(depth - 2): layers.append(nn.Linear(width, width)) layers.append(nn.Linear(width, output_dim)) model = nn.Sequential(*layers) for m in model.modules(): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, mode=mode, nonlinearity='relu') nn.init.zeros_(m.bias) return model x = torch.randn(64, input_dim) results = {} for mode in ['fan_in', 'fan_out']: model = build_model(mode) model.eval() with torch.no_grad(): # 前向传播 out = model(x) fwd_std = out.std().item() # 反向传播 model.train() model.zero_grad() loss = model(x).sum() loss.backward() # 第一层梯度的标准差 first_layer = list(model.modules())[1] bwd_std = first_layer.weight.grad.std().item() results[mode] = {'forward_std': fwd_std, 'backward_std': bwd_std} print(f"mode={mode}: forward_std={fwd_std:.4f}, backward_std={bwd_std:.6f}") return results # 运行对比 compare_init_modes(depth=30, width=256)

方案四：卷积层的初始化

import torch import torch.nn as nn class ConvNet(nn.Module): def __init__(self, num_classes=10): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d((1, 1)), ) self.classifier = nn.Linear(128, num_classes) self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): # 卷积层使用 fan_in + relu # fan_in = in_channels * kernel_size * kernel_size nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) elif isinstance(m, nn.Linear): # 最后一层分类器可以用 fan_in + relu 或 fan_in + linear nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') nn.init.zeros_(m.bias) def forward(self, x): x = self.features(x) x = x.flatten(1) x = self.classifier(x) return x

完整修复代码

""" Kaiming 初始化完整指南 包含: mode 选择、nonlinearity 匹配、效果对比 """ import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Optional, str class KaimingInitializer: """Kaiming 初始化管理器""" # 激活函数到 nonlinearity 的映射 ACTIVATION_MAP = { 'relu': ('relu', 0), 'leaky_relu': ('leaky_relu', 0.01), 'leaky_relu_0.2': ('leaky_relu', 0.2), 'prelu': ('leaky_relu', 0.25), 'elu': ('relu', 0), 'gelu': ('relu', 0), 'selu': ('linear', 0), 'tanh': ('tanh', 0), 'sigmoid': ('sigmoid', 0), 'linear': ('linear', 0), 'softmax': ('linear', 0), } @classmethod def init_layer(cls, layer: nn.Module, activation: str = 'relu', mode: str = 'fan_in'): """ 根据层类型和激活函数进行正确的 Kaiming 初始化 参数: layer: 要初始化的层 (Linear, Conv2d, etc.) activation: 实际使用的激活函数名称 mode: 'fan_in' 或 'fan_out' """ nonlinearity, a = cls.ACTIVATION_MAP.get(activation, ('relu', 0)) if isinstance(layer, nn.Linear): nn.init.kaiming_normal_( layer.weight, a=a, mode=mode, nonlinearity=nonlinearity ) if layer.bias is not None: nn.init.zeros_(layer.bias) elif isinstance(layer, nn.Conv2d): nn.init.kaiming_normal_( layer.weight, a=a, mode=mode, nonlinearity=nonlinearity ) if layer.bias is not None: nn.init.zeros_(layer.bias) elif isinstance(layer, nn.BatchNorm2d) or isinstance(layer, nn.BatchNorm1d): nn.init.ones_(layer.weight) nn.init.zeros_(layer.bias) @classmethod def init_model(cls, model: nn.Module, activation: str = 'relu', mode: str = 'fan_in', verbose: bool = True): """ 初始化整个模型 自动识别层类型并应用正确的初始化 """ stats = {'linear': 0, 'conv': 0, 'bn': 0, 'other': 0} for m in model.modules(): if isinstance(m, nn.Linear): cls.init_layer(m, activation, mode) stats['linear'] += 1 elif isinstance(m, nn.Conv2d): cls.init_layer(m, activation, mode) stats['conv'] += 1 elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)): cls.init_layer(m, activation, mode) stats['bn'] += 1 else: stats['other'] += 1 if verbose: print(f"Initialized: {stats}") print(f" mode={mode}, activation={activation}") # ============ 完整训练示例 ============ class ResidualBlock(nn.Module): def __init__(self, in_ch, out_ch, stride=1): super().__init__() self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(out_ch) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_ch) self.relu = nn.ReLU(inplace=True) self.downsample = None if stride != 1 or in_ch != out_ch: self.downsample = nn.Sequential( nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False), nn.BatchNorm2d(out_ch) ) def forward(self, x): identity = x out = self.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) if self.downsample: identity = self.downsample(x) return self.relu(out + identity) class ResNet(nn.Module): def __init__(self, num_blocks=[2, 2, 2, 2], num_classes=10): super().__init__() self.in_channels = 64 self.conv1 = nn.Conv2d(3, 64, 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.layer1 = self._make_layer(64, num_blocks[0], 1) self.layer2 = self._make_layer(128, num_blocks[1], 2) self.layer3 = self._make_layer(256, num_blocks[2], 2) self.layer4 = self._make_layer(512, num_blocks[3], 2) self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) self.fc = nn.Linear(512, num_classes) def _make_layer(self, out_ch, num_blocks, stride): layers = [ResidualBlock(self.in_channels, out_ch, stride)] self.in_channels = out_ch for _ in range(1, num_blocks): layers.append(ResidualBlock(out_ch, out_ch, 1)) return nn.Sequential(*layers) def forward(self, x): x = self.relu(self.bn1(self.conv1(x))) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.flatten(1) return self.fc(x) def main(): """完整使用示例""" print("=" * 60) print("Kaiming 初始化完整示例") print("=" * 60) # 1. 创建模型 model = ResNet(num_classes=10) # 2. 使用 fan_in 初始化（推荐） print("\n--- fan_in 初始化 ---") KaimingInitializer.init_model(model, activation='relu', mode='fan_in') # 3. 检查初始化效果 x = torch.randn(8, 3, 32, 32) model.eval() with torch.no_grad(): out = model(x) print(f"Input std: {x.std().item():.4f}") print(f"Output std: {out.std().item():.4f}") # 4. 对比 fan_out print("\n--- fan_out 初始化 ---") model2 = ResNet(num_classes=10) KaimingInitializer.init_model(model2, activation='relu', mode='fan_out') model2.eval() with torch.no_grad(): out2 = model2(x) print(f"Input std: {x.std().item():.4f}") print(f"Output std: {out2.std().item():.4f}") # 5. 梯度对比 print("\n--- 梯度对比 ---") for name, mdl in [('fan_in', model), ('fan_out', model2)]: mdl.train() mdl.zero_grad() loss = mdl(x).sum() loss.backward() grad_std = mdl.conv1.weight.grad.std().item() print(f"{name}: first layer grad std = {grad_std:.6f}") # 6. 训练对比 print("\n--- 训练对比 (5 epochs) ---") for name, mdl in [('fan_in', model), ('fan_out', model2)]: mdl.train() optimizer = torch.optim.SGD(mdl.parameters(), lr=0.01, momentum=0.9) criterion = nn.CrossEntropyLoss() for epoch in range(5): total_loss = 0 for _ in range(10): bx = torch.randn(16, 3, 32, 32) by = torch.randint(0, 10, (16,)) optimizer.zero_grad() out = mdl(bx) loss = criterion(out, by) loss.backward() optimizer.step() total_loss += loss.item() print(f" {name} epoch {epoch+1}: loss={total_loss/10:.4f}") if __name__ == "__main__": main()

常见陷阱与注意事项

1. fan_in vs fan_out 的选择

• 大多数情况用fan_in：保持前向传播信号方差稳定，是 PyTorch 的默认值。
• 深层网络反传困难时用fan_out：保持反向传播梯度方差稳定。
• 当 fan_in == fan_out 时无区别：如方阵形状的权重。

2. nonlinearity 必须与实际激活函数匹配

使用 ReLU 时设nonlinearity='relu'，使用 LeakyReLU 时设nonlinearity='leaky_relu'并正确设置a参数。

3. BatchNorm 层的初始化

BatchNorm 的 weight 初始化为 1，bias 初始化为 0，不要用 Kaiming 初始化。

4. 最后一层的初始化

分类任务的最后一层（分类头）可以考虑使用较小的初始化，以减小初始损失：

# 最后一层用较小的初始化 nn.init.kaiming_normal_(self.fc.weight, mode='fan_in', nonlinearity='linear') nn.init.zeros_(self.fc.bias)

5. 与 Xavier/Glorot 初始化的区别

• Xavier/Glorot：适用于 tanh、sigmoid 等饱和激活函数，使用fan_in + fan_out的平均。
• Kaiming/He：适用于 ReLU 等非饱和激活函数，使用fan_in或fan_out。

6. PyTorch 默认初始化

PyTorch 的nn.Linear默认使用kaiming_uniform_witha=sqrt(5)，这等价于一个较大的均匀分布初始化。手动使用kaiming_normal_通常能获得更好的训练效果。

总结

Kaiming 初始化的mode参数选择原则：

1. 默认使用fan_in：保持前向传播信号方差稳定，适用于绝大多数场景。
2. 深层网络考虑fan_out：当反向传播梯度消失时，使用fan_out保持梯度方差。
3. 匹配 nonlinearity：nonlinearity参数必须与实际使用的激活函数一致。
4. 正确设置 a 参数：使用 LeakyReLU 时，a应等于 negative_slope。
5. BatchNorm 特殊处理：BatchNorm 的 weight=1, bias=0，不用 Kaiming。
6. 卷积层同样适用：fan_in = in_channels * kernel_size^2。

理解fan_in和fan_out的数学含义，根据网络深度和激活函数选择正确的参数，是确保训练稳定性的关键。