极客笔记Pytorch 如何冻结PyTorch模型中的特定层
在本文中，我们将介绍如何在PyTorch模型中冻结特定层的方法。冻结某些层可以防止它们在训练过程中被更新，从而保持它们的权重不变。这在迁移学习、模型微调和特定任务中非常有用。
阅读更多：Pytorch 教程
什么是冻结层？
在深度学习中，神经网络模型通常由多个层组成，包括卷积层、全连接层等。当我们训练模型时，模型会根据输入数据和损失函数进行反向传播，并更新每个层的权重和偏置，以使模型能够更好地拟合训练数据。然而，在某些情况下，我们希望保持某些层的权重不变，以防止它们在训练过程中被更新。这就是冻结层的作用。
冻结层的方法
下面介绍两种常见的冻结层的方法：使用requires_grad属性和遍历模型的参数。
使用requires_grad属性
PyTorch中的张量（包括模型的权重和偏置）具有requires_grad属性，该属性决定是否计算其梯度以进行反向传播。默认情况下，所有张量的requires_grad属性为True，即会计算梯度。要冻结层，我们只需将层的参数的requires_grad属性设置为False即可。
以下示例演示了如何使用requires_grad属性来冻结特定层：
import torch
import torch.nn as nn
# 定义一个模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, 3)
self.conv2 = nn.Conv2d(64, 128, 3)
self.fc1 = nn.Linear(128 * 3 * 3, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return x
# 创建一个模型实例
model = MyModel()
# 冻结conv1层
model.conv1.requires_grad = False
# 输出模型中各层的requires_grad属性
for name, param in model.named_parameters():
print(f"{name}: {param.requires_grad}")
运行上述代码，我们可以看到conv1层的requires_grad属性已设置为False，而其他层的属性仍然为True。
遍历模型的参数
另一种冻结层的方法是遍历模型的参数，并将特定层的参数的requires_grad属性设置为False。这种方法可以更灵活地控制哪些层被冻结。
以下示例演示了如何遍历模型的参数并冻结特定层：
import torch
import torch.nn as nn
# 定义一个模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, 3)
self.conv2 = nn.Conv2d(64, 128, 3)
self.fc1 = nn.Linear(128 * 3 * 3, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return x
# 创建一个模型实例
model = MyModel()
# 冻结conv1层
for name, param in model.named_parameters():
if name.startswith('conv1'):
param.requires_grad = False
# 输出模型中各层的requires_grad属性
for name, param in model.named_parameters():
print(f"{name}: {param.requires_grad}")
运行上述代码，我们可以看到conv1层的requires_grad属性已设置为False，而其他层的属性仍然为True。
如何验证层是否已冻结？
为了验证我们是否成功地冻结了特定层，我们可以打印模型的反向传播过程中各层参数的梯度。如果冻结的层的梯度始终为None，则表示该层已成功冻结。
以下是一个示例：
import torch
import torch.nn as nn
# 定义一个模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.conv1 = nn.Conv2d(3, 64, 3)
self.conv2 = nn.Conv2d(64, 128, 3)
self.fc1 = nn.Linear(128 * 3 * 3, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = self.fc2(x)
return x
# 创建一个模型实例
model = MyModel()
# 冻结conv1层
model.conv1.requires_grad = False
# 随机生成输入数据
input_data = torch.randn(10, 3, 32, 32)
# 运行前向传播
output = model(input_data)
# 计算损失函数
loss = output.mean()
# 反向传播
loss.backward()
# 打印各层参数的梯度
for name, param in model.named_parameters():
print(f"{name}: {param.grad}")
运行以上代码，我们可以看到conv1层的梯度为None，而其他层的梯度不为None，这表明conv1层已成功冻结。
总结
本文介绍了如何在PyTorch模型中冻结特定层的方法。我们可以通过设置参数的requires_grad属性为False或遍历模型的参数并设置特定层的requires_grad属性为False来实现冻结层。冻结层对于迁移学习、模型微调和特定任务非常有用，可以保持特定层的权重不变，从而避免其在训练过程中被更新。希望本文对您了解PyTorch中冻结层的方法有所帮助！