本教程深入探讨了在pytorch中冻结神经网络特定中间层参数的两种常见方法：torch.no_grad()上下文管理器和设置参数的requires_grad = False属性。文章通过代码示例详细阐述了两种方法的原理、效果及适用场景，并明确指出requires_grad = False是实现精确中间层冻结的推荐方案，同时提供了验证层是否被冻结的技巧，旨在帮助开发者准确控制模型训练过程中的参数更新。
在深度学习模型训练过程中，我们经常会遇到需要冻结模型中某些层（即不更新这些层的参数）而只训练其他层的场景，例如在迁移学习中冻结预训练模型的特征提取层，或者在多任务学习中只更新特定任务相关的层。本文将详细探讨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的自动求导引擎不需要为这些参数计算梯度。
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模型中各层的训练状态，从而实现更复杂的训练策略，例如微调预训练模型或进行部分模型的更新。
