PyTorch Lightning权重初始化：默认方法查询及自定义方案

一、怎么查当前用的权重初始化方法 #

• 先看PyTorch原生层的默认逻辑：PyTorch Lightning不会改动PyTorch内置层（比如Linear、Conv2d）的默认初始化，这些层的初始化规则就是PyTorch自带的。你可以直接打印层的权重值观察分布，或者查PyTorch官方文档里对应层的初始化说明。举个例子：

  import torch.nn as nn
  linear_layer = nn.Linear(10, 5)
  print(linear_layer.weight)  # 从权重分布能看出是默认的kaiming_uniform初始化
  

• 检查自己的LightningModule代码：如果你的模型里有手动写的初始化逻辑，直接看__init__或者setup方法里有没有对权重赋值的代码就行。
• 预训练模型的情况：如果用了预训练权重，初始化方式就是预训练时用的，去对应预训练模型的文档或源码里找就行。

二、自定义修改权重初始化方式 #

1. 给单个层单独设置初始化 #

在LightningModule的__init__方法里，创建完层之后直接手动初始化：

import pytorch_lightning as pl
import torch.nn as nn
import torch.nn.init as init

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
        # 用xavier_normal初始化权重，偏置设为0
        init.xavier_normal_(self.linear.weight)
        init.constant_(self.linear.bias, 0.0)

2. 批量给同类型层统一初始化 #

遍历模型里的所有层，对指定类型的层批量设置初始化，比如在__init__里做：

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 20)
        self.linear2 = nn.Linear(20, 5)
        self.conv = nn.Conv2d(3, 16, 3)
        
        # 给所有Linear层用kaiming_normal初始化，偏置设0；Conv层用xavier_uniform
        for m in self.modules():
            if isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.Conv2d):
                init.xavier_uniform_(m.weight)

3. 借助Lightning钩子方法动态初始化 #

如果需要在训练/验证开始前动态调整初始化，可以用setup钩子（初始化只做一次的话，还是推荐在__init__里处理）：

class MyModel(pl.LightningModule):
    def setup(self, stage):
        # 在训练或验证阶段开始前初始化权重
        for m in self.modules():
            if isinstance(m, nn.Linear):
                init.normal_(m.weight, mean=0, std=0.01)

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