How does pytorch lock the parameters of the model, and you can customize how many layers you lock

In PyTorch, the parameters of some layers can be fixed by setting requires_gradthe attribute , thus locking the parameters of the model. At the same time, the number of layers to be locked can be selected according to specific needs.

The following is a step-by-step code example for locking model parameters and customizing the number of locked layers:

1. define model

First, we need to define the neural network model. For example, we can define a model with 4 convolutional layers and 3 fully connected layers, and set the parameters of the first two convolutional layers to not need to be updated (i.e. locked). It is assumed here that the input image size is 28x28, and the classification task is binary classification.

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

2. Freeze the number of layers and train the model

Next, we need to set requires_grad to False for the layers we want to freeze, and set requires_grad to True for the layers we don't want to freeze. Then, redefine the optimizer, loss function, etc., and train.

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()

In the above code, we first define an optimizer optimizerand loss function criterion. Then, set the requires_grad of the parameters that need to be frozen to False, and then use optim.SGD()the function to filter out the parameters whose requires_grad is True to update. Finally, model training is performed.

Through the above steps, we can lock the parameters of the PyTorch model and customize the number of layers that need to be locked.