通过torch.nn.init更改模型初始化参数

比如模型如下：

class Net(nn.Module):
    # 初始化定义网络的结构：也就是定义网络的层
    def __init__(self):
        super(Net,self).__init__()
        self.layer1 = nn.Sequential(            
            nn.Conv2d(3,6,kernel_size=5,stride=1,padding=0),
            # 激活函数
            nn.ReLU(),
            # kernel_size：pooling size，stride：down-sample
            nn.MaxPool2d(kernel_size=2,stride=2))
        self.layer2 = nn.Sequential(            
            nn.Conv2d(6,16,kernel_size=5,stride=1,padding=0),
            # 激活函数
            nn.ReLU(),
            # kernel_size：pooling size，stride：down-sample
            nn.MaxPool2d(kernel_size=2,stride=2))

        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        # x = x.reshape(x.size(0),-1)
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

model= Net().to(device)

在外部修改：

w1 = torch.empty(model.layer1[0].weight.shape) #according to the layer shape to create tensor

nn.init.kaiming_uniform_(w1, mode='fan_in', nonlinearity='relu')

model.layer1[0].weight.data.copy_(w1)

通过这种方式，模型的第一个卷积层的参数就被修改了。

这里用到的是He Kaiming的uniform的初始化方式，还有normal，以及很多其他人的方式也可以用，比如Xavier initialization， 详情可见这个页面， pytorch提供了一系列初始化参数的方法，可以对还没开始训练的模型进行参数修改。

当然也可以放在模型内部就修改。

class Net(nn.Module):
    # 初始化定义网络的结构：也就是定义网络的层
    def __init__(self):
        super(Net,self).__init__()
        self.layer1 = nn.Sequential(            
            nn.Conv2d(3,6,kernel_size=5,stride=1,padding=0),
            # 激活函数
            nn.ReLU(),
            # kernel_size：pooling size，stride：down-sample
            nn.MaxPool2d(kernel_size=2,stride=2))
        self.layer2 = nn.Sequential(            
            nn.Conv2d(6,16,kernel_size=5,stride=1,padding=0),
            # 激活函数
            nn.ReLU(),
            # kernel_size：pooling size，stride：down-sample
            nn.MaxPool2d(kernel_size=2,stride=2))
        w1 = torch.empty(self.layer1[0].weight.shape) #according to the layer shape to create tensor
        
        nn.init.kaiming_normal_(w1, mode='fan_in', nonlinearity='relu')
        
        self.layer1[0].weight.data.copy_(w1)

        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        # x = x.reshape(x.size(0),-1)
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

model= Net().to(device)