在深度学习模型训练过程中，我们经常会遇到需要冻结模型中某些层（即不更新这些层的参数）而只训练其他层的场景，例如在迁移学习中冻结预训练模型的特征提取层，或者在多任务学习中只更新特定任务相关的层。本文将详细探讨pytorch中实现这一目标的方法。
理解参数冻结的原理
在PyTorch中，参数更新是通过反向传播计算梯度并由优化器应用到参数上的。冻结一个层意味着阻止其参数参与梯度计算和随后的更新。这通常通过控制参数的requires_grad属性来实现。当requires_grad为False时，PyTorch的自动求导引擎将不会为该参数计算梯度，从而阻止其被优化器更新。
方法一：使用 torch.no_grad() 上下文管理器
torch.no_grad()是一个上下文管理器，它会禁用在其作用域内所有操作的梯度计算。这意味着，任何在with torch.no_grad():块中执行的操作，都不会构建计算图，也不会跟踪梯度。
让我们通过一个简单的三层线性网络为例来演示：
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.lin0 = nn.Linear(1, 2)
self.lin1 = nn.Linear(2, 2)
self.lin2 = nn.Linear(2, 10)
def forward_with_no_grad(self, x):
x = self.lin0(x)
with torch.no_grad():
x = self.lin1(x) # 尝试冻结lin1
x = self.lin2(x)
return x
# 实例化模型
model_no_grad = SimpleModel()
# 记录初始参数
initial_lin0_weight = model_no_grad.lin0.weight.clone()
initial_lin1_weight = model_no_grad.lin1.weight.clone()
initial_lin2_weight = model_no_grad.lin2.weight.clone()
# 模拟训练步骤
input_data = torch.randn(1, 1)
target = torch.randint(0, 10, (1,))
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model_no_grad.parameters(), lr=0.01)
print("--- 使用 torch.no_grad() 冻结中间层 ---")
print("初始 lin0 权重:\n", initial_lin0_weight)
print("初始 lin1 权重:\n", initial_lin1_weight)
print("初始 lin2 权重:\n", initial_lin2_weight)
# 前向传播与反向传播
output = model_no_grad.forward_with_no_grad(input_data)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 检查参数变化
print("\n训练后 lin0 权重:\n", model_no_grad.lin0.weight)
print("训练后 lin1 权重:\n", model_no_grad.lin1.weight)
print("训练后 lin2 权重:\n", model_no_grad.lin2.weight)
# 验证是否冻结
print("\nlin0 权重是否变化:", not torch.equal(initial_lin0_weight, model_no_grad.lin0.weight))
print("lin1 权重是否变化:", not torch.equal(initial_lin1_weight, model_no_grad.lin1.weight))
print("lin2 权重是否变化:", not torch.equal(initial_lin2_weight, model_no_grad.lin2.weight))分析 torch.no_grad() 的效果： 上述代码运行后会发现，lin0和lin1的参数都没有更新，而只有lin2的参数发生了变化。 这是因为当lin1的操作在torch.no_grad()块中执行时，其输出张量x（来自lin1）的grad_fn属性将为None，这意味着从lin1往前的计算图被截断了。因此，尽管lin2的梯度可以正常计算并回传到lin1的输出，但由于lin1的操作没有梯度跟踪，导致无法计算lin1自身的梯度，也无法将梯度继续回传到lin0。最终结果是，lin0和lin1的参数都不会得到更新。
结论： torch.no_grad() 适用于冻结整个模型或模型的一部分，使其在推理阶段不消耗内存来存储梯度信息，或者在训练时完全禁用某些部分的梯度更新。但它不适合精确地冻结中间层而允许其上游层更新的场景。
方法二：设置 requires_grad = False
这是在PyTorch中实现精确层冻结的推荐方法。通过将特定层的参数的requires_grad属性设置为False，我们可以明确告诉PyTorch的自动求导引擎不需要为这些参数计算梯度。
本书是全面讲述PHP与MySQL的经典之作，书中不但全面介绍了两种技术的核心特性，还讲解了如何高效地结合这两种技术构建健壮的数据驱动的应用程序。本书涵盖了两种技术新版本中出现的最新特性，书中大量实际的示例和深入的分析均来自于作者在这方面多年的专业经验，可用于解决开发者在实际中所面临的各种挑战。 本书内容全面深入，适合各层次PHP和MySQL开发人员阅读，既是优秀的学习教程，也可用作参考手册。
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个简单的模型
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.lin0 = nn.Linear(1, 2)
self.lin1 = nn.Linear(2, 2)
self.lin2 = nn.Linear(2, 10)
def forward(self, x):
x = self.lin0(x)
x = self.lin1(x)
x = self.lin2(x)
return x
# 实例化模型
model_requires_grad = SimpleModel()
# 冻结lin1层的参数
model_requires_grad.lin1.weight.requires_grad = False
model_requires_grad.lin1.bias.requires_grad = False
# 记录初始参数
initial_lin0_weight_rg = model_requires_grad.lin0.weight.clone()
initial_lin1_weight_rg = model_requires_grad.lin1.weight.clone()
initial_lin2_weight_rg = model_requires_grad.lin2.weight.clone()
# 注意：优化器只应传入 requires_grad 为 True 的参数
optimizer_rg = optim.SGD(filter(lambda p: p.requires_grad, model_requires_grad.parameters()), lr=0.01)
# 模拟训练步骤
input_data = torch.randn(1, 1)
target = torch.randint(0, 10, (1,))
criterion = nn.CrossEntropyLoss()
print("\n--- 使用 requires_grad = False 冻结中间层 ---")
print("初始 lin0 权重:\n", initial_lin0_weight_rg)
print("初始 lin1 权重:\n", initial_lin1_weight_rg)
print("初始 lin2 权重:\n", initial_lin2_weight_rg)
# 前向传播与反向传播
output = model_requires_grad(input_data)
loss = criterion(output, target)
optimizer_rg.zero_grad()
loss.backward()
optimizer_rg.step()
# 检查参数变化
print("\n训练后 lin0 权重:\n", model_requires_grad.lin0.weight)
print("训练后 lin1 权重:\n", model_requires_grad.lin1.weight)
print("训练后 lin2 权重:\n", model_requires_grad.lin2.weight)
# 验证是否冻结
print("\nlin0 权重是否变化:", not torch.equal(initial_lin0_weight_rg, model_requires_grad.lin0.weight))
print("lin1 权重是否变化:", not torch.equal(initial_lin1_weight_rg, model_requires_grad.lin1.weight))
print("lin2 权重是否变化:", not torch.equal(initial_lin2_weight_rg, model_requires_grad.lin2.weight))分析 requires_grad = False 的效果： 运行上述代码后，你会发现lin0和lin2的参数都得到了更新，而只有lin1的参数保持不变。 这是因为：
- lin1.weight.requires_grad = False和lin1.bias.requires_grad = False明确地告诉PyTorch不要为这些参数计算梯度。
- 在反向传播时，尽管梯度会流经lin1，但由于lin1的参数被标记为不需要梯度，PyTorch会跳过其梯度计算，并继续将梯度回传到lin0。
- 优化器在初始化时，通过filter(lambda p: p.requires_grad, model_requires_grad.parameters())确保它只接收那些requires_grad=True的参数进行更新。
结论： requires_grad = False 是实现精确冻结模型中特定层（包括中间层）的正确且推荐的方法。它允许梯度流经被冻结的层，但不会更新该层自身的参数，同时能将梯度正确地传递给更上游的层。
验证层是否被冻结
在实际操作中，可以通过以下几种方式来验证层是否成功被冻结：
检查 param.requires_grad 属性： 在设置后，可以打印出model.lin1.weight.requires_grad来确认其是否为False。
检查 param.grad 属性： 在执行loss.backward()之后，检查被冻结层的参数（例如model.lin1.weight.grad）是否为None。如果为None，则表示没有为该参数计算梯度。
检查参数值是否变化： 在训练循环开始前记录参数的初始值，经过一个或多个训练步骤后，再次检查这些参数的值。如果参数值未发生变化，则说明该层已被冻结。这正是本文示例代码中采用的方法。
总结与最佳实践
- 精确冻结中间层： 始终使用设置参数的requires_grad = False属性来冻结模型中的特定层。
- 优化器初始化： 当冻结部分层时，务必在初始化优化器时，只将那些requires_grad = True的参数传递给优化器。例如：optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01)。
- torch.no_grad() 的适用场景： torch.no_grad() 主要用于推理阶段，或者在训练过程中完全禁用某一部分的梯度计算，它会截断计算图，不适合需要梯度回传到上游层的场景。
- 模型状态： 冻结层与model.train()和model.eval()没有直接冲突。model.eval()主要影响nn.BatchNorm和nn.Dropout等层在训练和评估模式下的行为，而requires_grad控制的是参数是否更新。
通过理解和正确应用requires_grad = False，开发者可以灵活地控制PyTorch模型中各层的训练状态，从而实现更复杂的训练策略，例如微调预训练模型或进行部分模型的更新。
