pytorch自定义初始化权重的方法

PyTorch是一个流行的深度学习框架，它提供了许多内置的初始化权重方法。但是，有时候我们需要自定义初始化权重方法来更好地适应我们的模型。在本攻略中，我们将介绍如何自定义初始化权重方法。

方法1：使用nn.Module的apply()函数

我们可以使用nn.Module的apply()函数来自定义初始化权重方法。apply()函数可以递归地遍历整个模型，并对每个子模块应用指定的函数。以下是一个示例代码，演示了如何使用apply()函数自定义初始化权重方法：

import torch.nn as nn

# 自定义初始化权重方法
def init_weights(m):
    if isinstance(m, nn.Conv2d):
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.constant_(m.bias, 0)

# 定义CNN模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 实例化CNN模型并应用自定义初始化权重方法
net = Net()
net.apply(init_weights)

在上面的代码中，我们首先定义了一个init_weights()函数，该函数接受一个子模块作为参数，并使用nn.init.normal_()函数和nn.init.constant_()函数来初始化权重和偏置。然后，我们定义了一个Net类，该类继承自nn.Module类，并定义了CNN模型的各个层。在Net类的构造函数中，我们没有指定任何初始化权重方法。最后，我们实例化CNN模型，并使用apply()函数将自定义初始化权重方法应用于整个模型。

方法2：使用nn.Module的__init__()函数

我们还可以在nn.Module的__init__()函数中自定义初始化权重方法。在__init__()函数中，我们可以使用nn.init.normal_()函数和nn.init.constant_()函数来初始化权重和偏置。以下是一个示例代码，演示了如何在__init__()函数中自定义初始化权重方法：

import torch.nn as nn

# 定义CNN模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        nn.init.normal_(self.conv1.weight, mean=0, std=0.01)
        nn.init.constant_(self.conv1.bias, 0)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        nn.init.normal_(self.conv2.weight, mean=0, std=0.01)
        nn.init.constant_(self.conv2.bias, 0)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        nn.init.normal_(self.fc1.weight, mean=0, std=0.01)
        nn.init.constant_(self.fc1.bias, 0)
        self.fc2 = nn.Linear(120, 84)
        nn.init.normal_(self.fc2.weight, mean=0, std=0.01)
        nn.init.constant_(self.fc2.bias, 0)
        self.fc3 = nn.Linear(84, 10)
        nn.init.normal_(self.fc3.weight, mean=0, std=0.01)
        nn.init.constant_(self.fc3.bias, 0)

    def forward(self, x):
        x = self.pool(nn.functional.relu(self.conv1(x)))
        x = self.pool(nn.functional.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 实例化CNN模型
net = Net()

在上面的代码中，我们首先定义了一个Net类，该类继承自nn.Module类，并定义了CNN模型的各个层。在Net类的构造函数中，我们使用nn.init.normal_()函数和nn.init.constant_()函数来初始化权重和偏置。最后，我们实例化CNN模型。