如何使用kaiming_uniform_初始化torch.nn.Transformer权重

核心结论 #

不需要自定义编码器、解码器类，完全可以基于PyTorch内置的nn.Transformer模块直接完成Kaiming均匀初始化。

实现原理 #

nn.Transformer内部所有可训练的线性投影层（包括多头注意力的Q/K/V投影、输出投影、前馈网络的两层线性层）都以标准nn.Linear子模块的形式存在，你可以通过nn.Module自带的.modules()方法递归遍历模型所有层级的子模块，筛选出需要初始化的层统一应用torch.nn.init.kaiming_uniform_即可，不需要重写编码器、解码器的内部逻辑。

具体修改方案 #

你现有代码已经对词嵌入层做了Kaiming初始化，但init_weights方法仅覆盖了最终输出的generator层，没有覆盖内置Transformer的内部参数，同时存在缺少Tensor类型导入的小问题，按以下步骤修改即可：

1. 在文件开头补充缺失的导入：

   from torch import Tensor
   

2. 重写init_weights方法，遍历所有子模块完成初始化，注意不要对LayerNorm层错误应用Kaiming初始化（LayerNorm默认权重初始化为1、偏置初始化为0即可，强行用Kaiming初始化会破坏训练稳定性）：

   def init_weights(self):
       # 初始化输出生成层
       nn.init.kaiming_uniform_(self.generator.weight, a=math.sqrt(5))
       if self.generator.bias is not None:
           fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.generator.weight)
           bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
           nn.init.uniform_(self.generator.bias, -bound, bound)
       
       # 遍历内置Transformer的所有子模块初始化
       for module in self.transformer.modules():
           if isinstance(module, nn.Linear):
               # 所有线性层权重用Kaiming均匀初始化
               nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
               if module.bias is not None:
                   fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
                   bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
                   nn.init.uniform_(module.bias, -bound, bound)
           elif isinstance(module, nn.LayerNorm):
               # LayerNorm保持标准初始化
               nn.init.ones_(module.weight)
               nn.init.zeros_(module.bias)
   

3. 实例化模型后手动调用初始化方法即可生效：

   model = Transformer(
       src_vocab_size=your_src_vocab_size,
       tgt_vocab_size=your_tgt_vocab_size
       # 其余超参数按你的需求传入
   )
   model.init_weights()
   

注意事项 #

• 如果你将Transformer的激活函数换成GELU，不需要调整Kaiming初始化的核心逻辑，保持现有参数即可；如果用LeakyReLU则需要对应修改kaiming_uniform_的nonlinearity和a参数。
• Dropout层、位置编码的注册缓冲区没有需要训练的参数，不需要做初始化处理。

内容的提问来源于stack exchange，提问作者R00