pytorch初始化权重方法总结

import torch.nn as nn
import torch

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 3, (3, 3), stride=(1, 1), padding=1)
        self.bn1 = nn.BatchNorm2d(3)
        self.relu = nn.ReLU()
    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))

        return x
# init.uniform  a, b均值分布的上下限
class model_param_init(nn.Module):
    def __init__(self, model):
        super().__init__()
        assert isinstance(model, nn.Module), 'model not a class nn.Module'
        self.net = model
        self.initParam()
    def initParam(self):
        for param in self.net.parameters():
            # nn.init.zeros_(param)
            # nn.init.ones_(param)
            # nn.init.normal_(param, mean=0, std=1)
            # nn.init.uniform_(param, a=0, b=1)
            # nn.init.constant_(param, val=1)   # 将所有权重初始化为1
            # nn.init.eye_(param)  # 只能将二维的tensor初始化为单位矩阵
            # nn.init.xavier_uniform_(param, gain=1)  # Glorot初始化  得到的张量是从-a——a中采用的
            # nn.init.xavier_normal_(param, gain=1)   # 得到的张量是从0-std采样的
            nn.init.kaiming_normal_(param, a=0, mode='fan_in', nonlinearity='leaky_relu') # he初始化方法
            # nn.init.kaiming_uniform_(param)

if __name__ == '__main__':
    net = Net()
    net = model_param_init(net)
    for param in net.parameters():
        print(param)
    # 按照参数更改权重
    for name, param in net.named_parameters():
        if name == net.conv1.weight:  # 指定更改某层的权重值
            nn.init.dirac_(param, groups=1) # 保留卷积层通道的值
        print(param)

这里可以将model_param_init封装成一个函数，从而在更改初始化方法时更加方便。