pytorch如何锁住模型的参数,并且可以自定义锁住几层

在 PyTorch 中,可以通过设置 requires_grad 属性来固定某些层的参数,从而锁住模型的参数。同时,可以根据具体需求选择需要锁住的层数。

以下是一步一步实现锁定模型参数并自定义锁住层数的代码示例:

  1. 定义模型

首先,我们需要定义神经网络模型。例如,我们可以定义一个包含 4 个卷积层和 3 个全连接层的模型,并将前两个卷积层的参数设为不需要更新(即锁住)。这里假设输入的图像大小为 28x28,分类任务为二分类。

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        # 卷积层1-4,输出通道数分别为32、64、128、128
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=(3, 3), padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=(3, 3), padding=1)
        self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=(3, 3), padding=1)
        self.conv4 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=(3, 3), padding=1)
        # 全连接层1-3,输出维度分别为64、32、2
        self.fc1 = nn.Linear(128 * 7 * 7, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 2)

    def forward(self, x):
        # 卷积层1-4,激活函数使用ReLU
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.relu(self.conv3(x))
        # 锁定第四个卷积层的参数
        with torch.no_grad():
            x = nn.functional.relu(self.conv4(x))
        # 池化层
        x = nn.functional.max_pool2d(x, kernel_size=2)
        # 展平
        x = x.view(-1, 128 * 7 * 7)
        # 全连接层1-3,激活函数使用ReLU
        x = nn.functional.relu(self.fc1(x))
        # 锁定全连接层1和2
        with torch.no_grad():
            x = nn.functional.relu(self.fc2(x))
        x = self.fc3(x)
        return x
  1. 冻结层数并训练模型

接下来,我们需要将要冻结的层的 requires_grad 设置为 False,并将不需要冻结的层的 requires_grad 设置为 True。然后,重新定义优化器、损失函数等内容,并进行训练。

import torch.optim as optim

model = MyModel()
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01, momentum=0.9)
criterion = nn.CrossEntropyLoss()

# 冻结第3个和第4个卷积层、全连接层1和2的参数
for param in model.conv3.parameters():
    param.requires_grad = False
for param in model.conv4.parameters():
    param.requires_grad = False
for param in model.fc1.parameters():
    param.requires_grad = False
for param in model.fc2.parameters():
    param.requires_grad = False

# 训练模型
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

在上述代码中,我们首先定义了一个优化器 optimizer 和损失函数 criterion。然后,将需要冻结的参数的 requires_grad 设置为 False,然后通过 optim.SGD() 函数过滤出 requires_grad 为 True 的参数来更新。最后进行模型训练。

通过以上步骤,我们就可以锁住 PyTorch 模型的参数,并自定义需要锁住的层数了。

おすすめ

転載: blog.csdn.net/weixin_45277161/article/details/130744162