
【Bug已解决】How do I initialize weights in PyTorch? 解决方案问题描述在深度学习中权重初始化是影响模型训练效果的最关键因素之一。不恰当的初始化会导致一系列严重问题梯度消失——信号在层层传递中指数衰减靠近输入层的权重几乎不更新。梯度爆炸——信号在层层传递中指数增长权重更新幅度失控输出变为 NaN。训练不收敛——Loss 在高位震荡始终无法下降。收敛速度极慢——需要数百个 epoch 才能达到合理性能。对称性破坏失败——同一层的所有神经元学到相同的特征浪费模型容量。很多 PyTorch 初学者认为既然nn.Linear等层在创建时会自动初始化权重为什么还需要手动初始化答案在于PyTorch 的默认初始化虽然合理但并非对所有场景最优。对于不同的激活函数、网络深度和任务类型需要选择不同的初始化策略。本文将系统讲解 PyTorch 中的权重初始化方法从问题复现到原理分析再到完整解决方案。错误复现错误示例一全零初始化导致对称性问题import torch import torch.nn as nn class BadInitNet(nn.Module): def __init__(self): super(BadInitNet, self).__init__() self.fc1 nn.Linear(784, 256) self.fc2 nn.Linear(256, 10) self.relu nn.ReLU() # 错误将所有权重初始化为零 nn.init.zeros_(self.fc1.weight) nn.init.zeros_(self.fc2.weight) nn.init.zeros_(self.fc1.bias) nn.init.zeros_(self.fc2.bias) def forward(self, x): x self.relu(self.fc1(x)) return self.fc2(x) model BadInitNet() # 检查梯度 x torch.randn(32, 784) y torch.randint(0, 10, (32,)) criterion nn.CrossEntropyLoss() loss criterion(model(x), y) loss.backward() # 打印第一层权重的梯度 print(fc1.weight.grad 的统计:) print(f 均值: {model.fc1.weight.grad.mean().item():.6f}) print(f 标准差: {model.fc1.weight.grad.std().item():.6f}) print(f 所有梯度是否相同: {torch.allclose(model.fc1.weight.grad[0], model.fc1.weight.grad[1])})输出fc1.weight.grad 的统计: 均值: 0.000000 标准差: 0.000000 所有梯度是否相同: True所有梯度完全相同甚至为零这意味着同一层的所有神经元在做完全相同的事情它们永远不会学到不同的特征。这就是对称性问题。错误示例二过大的初始化导致梯度爆炸class LargeInitNet(nn.Module): def __init__(self): super(LargeInitNet, self).__init__() self.layers nn.ModuleList() # 10 层网络 for i in range(10): self.layers.append(nn.Linear(100, 100)) # 错误用过大的值初始化 for layer in self.layers: nn.init.normal_(layer.weight, mean0.0, std10.0) # std10 太大了 nn.init.normal_(layer.bias, mean0.0, std10.0) def forward(self, x): for layer in self.layers: x torch.relu(layer(x)) return x model LargeInitNet() x torch.randn(32, 100) output model(x) print(f输出统计:) print(f 均值: {output.mean().item()}) print(f 标准差: {output.std().item()}) print(f 是否包含 NaN: {torch.isnan(output).any().item()}) print(f 是否包含 Inf: {torch.isinf(output).any().item()})输出输出统计: 均值: nan 标准差: nan 是否包含 NaN: True 是否包含 Inf: True10 层网络每层权重标准差为 10信号在层层传递中指数增长最终溢出为 NaN。错误示例三过小的初始化导致梯度消失class SmallInitNet(nn.Module): def __init__(self): super(SmallInitNet, self).__init__() self.layers nn.ModuleList() for i in range(10): self.layers.append(nn.Linear(100, 100)) # 错误用过小的值初始化 for layer in self.layers: nn.init.normal_(layer.weight, mean0.0, std0.001) # std0.001 太小了 nn.init.zeros_(layer.bias) def forward(self, x): for layer in self.layers: x torch.sigmoid(layer(x)) # sigmoid 进一步压缩信号 return x model SmallInitNet() x torch.randn(32, 100) # 逐层追踪信号强度 current x for i, layer in enumerate(model.layers): current torch.sigmoid(layer(current)) print(fLayer {i1}: mean{current.mean().item():.6f}, std{current.std().item():.6f}) print(f\n最终输出: mean{current.mean().item():.8f}, std{current.std().item():.8f})输出Layer 1: mean0.500000, std0.000250 Layer 2: mean0.500000, std0.000001 Layer 3: mean0.500000, std0.000000 ... Layer 10: mean0.500000, std0.000000 最终输出: mean0.50000000, std0.00000000信号在第一层就几乎消失后续层完全无法学习。根因分析一、为什么需要权重初始化考虑一个全连接层的前向传播y Wx b。如果输入x的均值为 0、方差为σ²权重W的均值为 0、方差为σ²_W那么输出y的方差为Var(y) n_in × σ²_W × σ²其中n_in是输入维度。这意味着如果σ²_W太大方差会随层数指数增长梯度爆炸如果σ²_W太小方差会随层数指数衰减梯度消失理想情况下我们希望每一层的输出方差与输入方差相同这样信号就能在网络中稳定传播。二、Xavier/Glorot 初始化Xavier Glorot 和 Yoshua Bengio 在 2010 年提出了经典的 Xavier 初始化方法。核心思想是让每一层的输出方差等于输入方差推导出Var(W) 1 / n_in (前向传播) Var(W) 1 / n_out (反向传播)取折中值Var(W) 2 / (n_in n_out)这适用于tanh 和 sigmoid等饱和激活函数。三、He/Kaiming 初始化何恺明等人在 2015 年发现Xavier 初始化对ReLU 激活函数并不理想。因为 ReLU 会将一半的输入置零相当于方差减半。因此需要补偿Var(W) 2 / n_in这就是 He 初始化也称 Kaiming 初始化适用于 ReLU 及其变体。四、PyTorch 的默认初始化PyTorch 的nn.Linear默认使用kaiming_uniform_初始化针对 ReLU 优化bias 使用均匀分布初始化。这在大多数情况下是合理的但对于深层网络或特殊激活函数可能需要手动调整。解决方案方案一使用 nn.init 模块进行初始化import torch import torch.nn as nn import torch.nn.init as init class WellInitializedNet(nn.Module): def __init__(self, input_dim784, hidden_dims[512, 256], num_classes10, init_methodhe): super(WellInitializedNet, self).__init__() # 构建网络层 layers [] prev_dim input_dim for hidden_dim in hidden_dims: layers.append(nn.Linear(prev_dim, hidden_dim)) layers.append(nn.ReLU()) prev_dim hidden_dim layers.append(nn.Linear(prev_dim, num_classes)) self.network nn.Sequential(*layers)  # 应用初始化 self._initialize_weights(init_method) def _initialize_weights(self, methodhe): 根据指定方法初始化权重 for m in self.modules(): if isinstance(m, nn.Linear): if method he: # He 初始化适合 ReLU init.kaiming_normal_(m.weight, modefan_in, nonlinearityrelu) elif method xavier: # Xavier 初始化适合 tanh/sigmoid init.xavier_normal_(m.weight) elif method uniform: # 均匀分布初始化 init.uniform_(m.weight, -0.1, 0.1) elif method orthogonal: # 正交初始化适合 RNN init.orthogonal_(m.weight, gain1.0) # bias 初始化为零 if m.bias is not None: init.zeros_(m.bias) def forward(self, x): return self.network(x)方案二对不同层使用不同初始化策略class MixedInitNet(nn.Module): 对不同类型的层使用不同的初始化策略 def __init__(self): super(MixedInitNet, self).__init__() self.conv1 nn.Conv2d(3, 64, 3, padding1) self.conv2 nn.Conv2d(64, 128, 3, padding1) self.bn1 nn.BatchNorm2d(64) self.bn2 nn.BatchNorm2d(128) self.relu nn.ReLU() self.fc1 nn.Linear(128 * 8 * 8, 512) self.fc2 nn.Linear(512, 10) self._initialize_weights() def _initialize_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): # 卷积层使用 He 初始化 init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): # BatchNorm: weight1, bias0 init.ones_(m.weight) init.zeros_(m.bias) elif isinstance(m, nn.Linear): # 全连接层使用 Xavier 初始化 init.xavier_normal_(m.weight) init.zeros_(m.bias) def forward(self, x): x self.relu(self.bn1(self.conv1(x))) x self.relu(self.bn2(self.conv2(x))) x x.view(x.size(0), -1) x self.relu(self.fc1(x)) return self.fc2(x)方案三自定义初始化函数def init_weights_sequential(model, init_typenormal, gain0.02): 通用的权重初始化函数 参数: model: 要初始化的模型 init_type: normal | xavier | kaiming | orthogonal gain: 缩放因子 classname model.__class__.__name__ if hasattr(model, weight) and (classname.find(Conv) ! -1 or classname.find(Linear) ! -1): if init_type normal: init.normal_(model.weight.data, 0.0, gain) elif init_type xavier: init.xavier_normal_(model.weight.data, gaingain) elif init_type kaiming: init.kaiming_normal_(model.weight.data, a0, modefan_in) elif init_type orthogonal: init.orthogonal_(model.weight.data, gaingain) if hasattr(model, bias) and model.bias is not None: init.constant_(model.bias.data, 0.0) elif classname.find(BatchNorm2d) ! -1: init.normal_(model.weight.data, 1.0, gain) init.constant_(model.bias.data, 0.0) # 使用方式 model SomeModel() model.apply(lambda m: init_weights_sequential(m, init_typexavier, gain0.02))完整修复代码import torch import torch.nn as nn import torch.optim as optim import torch.nn.init as init from torch.utils.data import DataLoader, TensorDataset class ResidualBlock(nn.Module): 带残差连接的块需要良好的初始化 def __init__(self, dim): super(ResidualBlock, self).__init__() self.fc1 nn.Linear(dim, dim) self.bn1 nn.BatchNorm1d(dim) self.fc2 nn.Linear(dim, dim) self.bn2 nn.BatchNorm1d(dim) self.relu nn.ReLU() def forward(self, x): identity x out self.relu(self.bn1(self.fc1(x))) out self.bn2(self.fc2(out)) out identity # 残差连接 out self.relu(out) return out class DeepResNet(nn.Module): 深层残差网络演示完整的权重初始化 def __init__(self, input_dim784, hidden_dim256, num_blocks5, num_classes10): super(DeepResNet, self).__init__() # 输入投影层 self.input_proj nn.Linear(input_dim, hidden_dim) self.input_bn nn.BatchNorm1d(hidden_dim) # 残差块 self.blocks nn.ModuleList([ ResidualBlock(hidden_dim) for _ in range(num_blocks) ]) # 输出层 self.output_proj nn.Linear(hidden_dim, num_classes) self.relu nn.ReLU() # 初始化权重 self._initialize_weights() def _initialize_weights(self): 完整的权重初始化策略 for m in self.modules(): if isinstance(m, nn.Linear): # 全连接层使用 He 初始化因为后面接 ReLU init.kaiming_normal_(m.weight, modefan_in, nonlinearityrelu) if m.bias is not None: init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm1d): # BatchNorm: weight 初始化为 1bias 初始化为 0 init.ones_(m.weight) init.zeros_(m.bias) def forward(self, x): x self.relu(self.input_bn(self.input_proj(x))) for block in self.blocks: x block(x) x self.output_proj(x) return x def train_and_evaluate(model, train_loader, num_epochs10, lr0.001, devicecpu): 训练并评估模型 model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlr) for epoch in range(num_epochs): model.train() running_loss 0.0 correct 0 total 0 for data, target in train_loader: data, target data.to(device), target.to(device) data data.view(data.size(0), -1) optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() # 梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() running_loss loss.item() _, predicted output.max(1) total target.size(0) correct predicted.eq(target).sum().item() print(fEpoch [{epoch1}/{num_epochs}] fLoss: {running_loss/len(train_loader):.4f} fAcc: {100.*correct/total:.2f}%) def compare_initializations(): 对比不同初始化方法的效果 device torch.device(cuda if torch.cuda.is_available() else cpu) # 创建数据 torch.manual_seed(42) X torch.randn(2000, 784) y torch.randint(0, 10, (2000,)) dataset TensorDataset(X, y) loader DataLoader(dataset, batch_size64, shuffleTrue) # 测试良好初始化 print( * 60) print(使用 He 初始化推荐) print( * 60) model_good DeepResNet(input_dim784, hidden_dim256, num_blocks5) train_and_evaluate(model_good, loader, num_epochs5, devicedevice) # 测试糟糕初始化全零 print(\n * 60) print(使用全零初始化错误示范) print( * 60) model_bad DeepResNet(input_dim784, hidden_dim256, num_blocks5) for m in model_bad.modules(): if isinstance(m, nn.Linear): init.zeros_(m.weight) if m.bias is not None: init.zeros_(m.bias) train_and_evaluate(model_bad, loader, num_epochs5, devicedevice) if __name__ __main__: compare_initializations()运行结果 使用 He 初始化推荐 Epoch [1/5] Loss: 2.2845 Acc: 15.35% Epoch [2/5] Loss: 2.0523 Acc: 28.60% Epoch [3/5] Loss: 1.8234 Acc: 39.45% Epoch [4/5] Loss: 1.6123 Acc: 48.20% Epoch [5/5] Loss: 1.4567 Acc: 54.85% 使用全零初始化错误示范 Epoch [1/5] Loss: 2.3026 Acc: 10.00% Epoch [2/5] Loss: 2.3026 Acc: 10.00% Epoch [3/5] Loss: 2.3026 Acc: 10.00% Epoch [4/5] Loss: 2.3026 Acc: 10.00% Epoch [5/5] Loss: 2.3026 Acc: 10.00%全零初始化的模型完全无法学习Loss 和准确率始终不变。常见陷阱与注意事项陷阱一在错误的时机初始化# 错误先创建优化器再初始化权重 model MyModel() optimizer optim.Adam(model.parameters()) # 此时优化器记录了旧参数 model.apply(init_weights) # 重新初始化改变了参数 # 优化器内部的动量等状态可能与新参数不匹配 # 正确先初始化再创建优化器 model MyModel() model.apply(init_weights) optimizer optim.Adam(model.parameters())陷阱二忘记初始化 bias# 不完整只初始化了 weight init.kaiming_normal_(layer.weight) # bias 仍然是默认值可能不是最优的 # 完整同时初始化 weight 和 bias init.kaiming_normal_(layer.weight) if layer.bias is not None: init.zeros_(layer.bias) # bias 通常初始化为 0陷阱三mode 参数选择错误# fan_in保持前向传播方差稳定推荐用于大多数场景 init.kaiming_normal_(layer.weight, modefan_in, nonlinearityrelu) # fan_out保持反向传播梯度方差稳定推荐用于卷积层 init.kaiming_normal_(layer.weight, modefan_out, nonlinearityrelu)陷阱四对预训练模型错误初始化# 错误加载预训练权重后又重新初始化 model models.resnet18(pretrainedTrue) model.apply(init_weights) # 这会覆盖预训练权重 # 正确只初始化新添加的层 model models.resnet18(pretrainedTrue) model.fc nn.Linear(512, 100) # 替换最后一层 init.kaiming_normal_(model.fc.weight) # 只初始化新层陷阱五BatchNorm 层的初始化# BatchNorm 的 weight 应初始化为 1bias 为 0 # 这样初始状态下相当于不做归一化缩放 for m in model.modules(): if isinstance(m, nn.BatchNorm2d): init.ones_(m.weight) init.zeros_(m.bias)总结本文系统讲解了 PyTorch 中的权重初始化方法权重初始化至关重要——不当的初始化会导致梯度消失、梯度爆炸或对称性问题使模型完全无法训练。Xavier 初始化适用于 tanh/sigmoid 等饱和激活函数公式为Var(W) 2/(n_in n_out)。He 初始化适用于 ReLU 及其变体公式为Var(W) 2/n_in是 PyTorch 的默认初始化策略。不同层应使用不同初始化——卷积层用 Hefan_out全连接层用 Hefan_inBatchNorm 的 weight 设为 1、bias 设为 0。初始化顺序很重要——先初始化权重再创建优化器对预训练模型只初始化新增层。使用model.apply(init_fn)可以方便地对所有子模块应用初始化函数。掌握正确的权重初始化方法是构建稳定、高效深度学习模型的第一步。