1. 引言
在深度学习模型的训练过程中，我们经常需要冻结（即停止更新）网络中的某些层，而只训练其他层。这在迁移学习、特征提取或调试网络行为时尤为常见。pytorch提供了灵活的机制来控制参数的梯度计算和更新。然而，对于如何精确地冻结网络中的特定中间层，同时允许其前后的层正常训练，常常会引起混淆。本文将详细探讨两种常见的冻结策略，并分析它们的实际效果。
2. 冻结策略一：使用torch.no_grad()上下文管理器
torch.no_grad()是一个上下文管理器，它会禁用梯度计算。在其作用域内的所有操作都不会构建计算图，因此不会计算梯度。
2.1 机制解释
当在模型的forward方法中使用with torch.no_grad():包裹某个层时，该层在正向传播过程中将不会追踪其操作，这意味着不会为该层的参数计算梯度。然而，这种方法的副作用是，一旦梯度流在某个点被no_grad中断，反向传播将无法穿透这个点到达其之前的层。这意味着，不仅被包裹的层参数不会更新，其之前的所有层（如果它们依赖于no_grad块的输出）的参数也将无法接收到梯度，从而也无法更新。
2.2 示例代码
考虑一个简单的三层线性网络，我们尝试使用torch.no_grad()冻结中间层lin1：
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleModelNoGrad(nn.Module):
def __init__(self):
super(SimpleModelNoGrad, self).__init__()
self.lin0 = nn.Linear(1, 2)
self.lin1 = nn.Linear(2, 2)
self.lin2 = nn.Linear(2, 1) # 修改输出维度以简化示例
def forward(self, x):
x = self.lin0(x)
with torch.no_grad():
x = self.lin1(x) # 尝试冻结此层
x = self.lin2(x)
return x
# 实例化模型
model_nograd = SimpleModelNoGrad()
# 打印初始参数值
print("--- 使用 torch.no_grad() ---")
print("初始 lin0.weight:", model_nograd.lin0.weight.data)
print("初始 lin1.weight:", model_nograd.lin1.weight.data)
print("初始 lin2.weight:", model_nograd.lin2.weight.data)
# 训练设置
criterion = nn.MSELoss()
optimizer = optim.SGD(model_nograd.parameters(), lr=0.01)
# 模拟训练步骤
input_data = torch.randn(1, 1)
target_data = torch.randn(1, 1)
for i in range(5):
optimizer.zero_grad()
output = model_nograd(input_data)
loss = criterion(output, target_data)
loss.backward()
optimizer.step()
# 打印训练后的参数值
print("\n训练后 lin0.weight:", model_nograd.lin0.weight.data)
print("训练后 lin1.weight:", model_nograd.lin1.weight.data)
print("训练后 lin2.weight:", model_nograd.lin2.weight.data)
# 验证参数是否更新
# 观察发现，lin0和lin1的参数都不会更新，因为no_grad阻断了梯度流2.3 效果分析
如上述示例所示，当lin1被torch.no_grad()包裹时，其参数不会更新。更重要的是，由于梯度流被阻断，lin0的参数也无法接收到梯度，因此lin0的参数也不会更新。只有lin2的参数会更新，因为它在no_grad块之后。这表明torch.no_grad()是一种“一刀切”的冻结方式，它会阻止梯度流向其之前的所有层。
3. 冻结策略二：设置参数的requires_grad属性为False
更精确和推荐的冻结方法是直接设置目标层中所有参数的requires_grad属性为False。
3.1 机制解释
每个nn.Module的参数（如weight和bias）都有一个requires_grad属性，默认为True。当此属性设置为False时，PyTorch的自动求导引擎将不会为这些参数计算梯度，也不会将这些参数包含在优化器的参数列表中。这使得梯度可以自由地流过这些冻结的层，并继续传递给它们之前的层，从而允许之前的层正常更新。
3.2 示例代码
我们再次考虑一个三层线性网络，这次使用requires_grad = False来冻结中间层lin1：
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleModelRequiresGrad(nn.Module):
def __init__(self):
super(SimpleModelRequiresGrad, self).__init__()
self.lin0 = nn.Linear(1, 2)
self.lin1 = nn.Linear(2, 2)
self.lin2 = nn.Linear(2, 1)
def forward(self, x):
x = self.lin0(x)
x = self.lin1(x)
x = self.lin2(x)
return x
# 实例化模型
model_reqgrad = SimpleModelRequiresGrad()
# 冻结 lin1 的所有参数
for param in model_reqgrad.lin1.parameters():
param.requires_grad = False
# 打印初始参数值
print("\n--- 使用 requires_grad = False ---")
print("初始 lin0.weight:", model_reqgrad.lin0.weight.data)
print("初始 lin1.weight:", model_reqgrad.lin1.weight.data)
print("初始 lin2.weight:", model_reqgrad.lin2.weight.data)
# 验证 requires_grad 状态
print("\nlin0.weight.requires_grad:", model_reqgrad.lin0.weight.requires_grad)
print("lin1.weight.requires_grad:", model_reqgrad.lin1.weight.requires_grad)
print("lin2.weight.requires_grad:", model_reqgrad.lin2.weight.requires_grad)
# 训练设置
# 注意：优化器只应接收 requires_grad=True 的参数
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model_reqgrad.parameters()), lr=0.01)
criterion = nn.MSELoss()
# 模拟训练步骤
input_data = torch.randn(1, 1)
target_data = torch.randn(1, 1)
for i in range(5):
optimizer.zero_grad()
output = model_reqgrad(input_data)
loss = criterion(output, target_data)
loss.backward()
optimizer.step()
# 打印训练后的参数值
print("\n训练后 lin0.weight:", model_reqgrad.lin0.weight.data)
print("训练后 lin1.weight:", model_reqgrad.lin1.weight.data)
print("训练后 lin2.weight:", model_reqgrad.lin2.weight.data)
# 验证参数是否更新
# 观察发现，lin0和lin2的参数会更新，而lin1的参数不会更新3.3 效果分析
通过设置model_reqgrad.lin1中所有参数的requires_grad = False，我们成功地冻结了lin1。此时，PyTorch的自动求导机制会为lin0和lin2的参数计算梯度，并允许优化器只更新那些requires_grad为True的参数。因此，lin0和lin2的参数会正常更新，而lin1的参数则保持不变。这种方法实现了对特定层参数的精确控制。
4. 总结与注意事项
-
torch.no_grad()：
- 优点：简单易用，适用于在推理阶段或完全禁用整个模型或某个子图的梯度计算。
- 缺点：当用于冻结中间层时，会阻断反向传播流，导致其之前的层也无法更新。不适用于需要选择性冻结中间层而同时训练其前置层的场景。
-
设置requires_grad = False：
- 优点：实现精确的层冻结，只冻结目标层的参数，而允许梯度流通过并更新其前置层。这是在迁移学习或需要精细控制训练过程时的推荐方法。
-
注意事项：
- 确保在定义优化器时，只将requires_grad = True的参数传递给优化器，例如：optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01)。否则，优化器可能会尝试更新那些不计算梯度的参数，导致不必要的计算或潜在错误。
- 冻结操作通常在模型实例化后、训练循环开始前进行。
- 可以通过打印param.requires_grad来验证参数的梯度计算状态，并通过比较训练前后的param.data来验证参数是否实际更新。
综上所述，当目标是选择性地冻结神经网络中的特定中间层，同时确保其他层能够正常训练时，设置参数的requires_grad属性为False是更准确和推荐的策略。
