在 PyTorch 中,可以通过设置 requires_grad
属性来固定某些层的参数,从而锁住模型的参数。同时,可以根据具体需求选择需要锁住的层数。
以下是一步一步实现锁定模型参数并自定义锁住层数的代码示例:
- 定义模型
首先,我们需要定义神经网络模型。例如,我们可以定义一个包含 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
- 冻结层数并训练模型
接下来,我们需要将要冻结的层的 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 模型的参数,并自定义需要锁住的层数了。