PyTorch中精确冻结中间层参数的策略与实践

本教程深入探讨了在PyTorch模型训练中冻结特定中间层参数的两种常见方法：使用torch.no_grad()上下文管理器和直接设置参数的requires_grad属性。通过实验对比，我们揭示了torch.no_grad()可能对上游层产生意外影响，而requires_grad = False是实现精确、选择性层冻结的推荐方案，这对于迁移学习和模型微调至关重要。

在深度学习模型训练中，冻结模型的部分层参数是一个常见的需求，尤其是在迁移学习、模型微调或实验特定层行为时。例如，我们可能希望在预训练模型的基础上，只训练顶层分类器，而保留底层特征提取器的参数不变。然而，如何正确且精确地冻结中间层，同时确保其他层的参数能够正常更新，是pytorch用户经常遇到的问题。本文将详细分析两种常用的方法，并通过代码示例阐明它们的行为差异。

理解参数冻结的需求

假设我们有一个由多个线性层组成的简单神经网络：lin0 -> lin1 -> lin2。我们的目标是冻结lin1层的参数，使其在训练过程中不发生更新，而lin0和lin2层的参数则应正常参与梯度计算和优化器更新。

方法一：使用 torch.no_grad() 上下文管理器

torch.no_grad() 是一个上下文管理器，用于禁用梯度计算。在它作用域内的所有计算都不会构建计算图，这意味着任何在此作用域内产生的张量都不会有grad_fn属性，从而无法进行反向传播。

代码示例：在 forward 方法中使用 torch.no_grad()

import torch
import torch.nn as nn

class SimpleModelWithNoGrad(nn.Module):
    def __init__(self):
        super(SimpleModelWithNoGrad, 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)
        # 在lin1的计算中使用torch.no_grad()
        with torch.no_grad():
            x = self.lin1(x)
        x = self.lin2(x)
        return x

# 实例化模型
model_no_grad = SimpleModelWithNoGrad()

# 打印初始参数的requires_grad属性
print("--- 使用 torch.no_grad() 时的初始requires_grad ---")
print(f"lin0.weight.requires_grad: {model_no_grad.lin0.weight.requires_grad}")
print(f"lin1.weight.requires_grad: {model_no_grad.lin1.weight.requires_grad}")
print(f"lin2.weight.requires_grad: {model_no_grad.lin2.weight.requires_grad}")

行为分析：

当我们尝试在forward方法中对lin1的计算使用with torch.no_grad():时，实验结果表明，不仅lin1的参数不会更新，连lin0的参数也未能更新。这是因为torch.no_grad()在执行lin1(x)时切断了从lin1到lin0的梯度流。由于lin1的输出不再追踪梯度，因此后续的反向传播无法将梯度信息传递回lin0，导致lin0的参数也无法得到更新。

结论： 这种方法并不适用于仅冻结中间层而让其上游层正常训练的场景。它更适用于在验证、推理阶段或模型中某个子模块确实不需要梯度计算时使用，以节省内存和计算。

方法二：设置 layer.requires_grad = False

PyTorch中的每个nn.Parameter（包括层的权重和偏置）都有一个requires_grad属性，默认为True。当这个属性设置为False时，PyTorch的反向传播机制会跳过这些参数，不计算它们的梯度，因此优化器也不会更新它们。

代码示例：设置 requires_grad = False

import torch
import torch.nn as nn

class SimpleModelWithRequiresGrad(nn.Module):
    def __init__(self):
        super(SimpleModelWithRequiresGrad, 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 = SimpleModelWithRequiresGrad()

# 冻结lin1的参数
# 确保对权重和偏置都进行设置
model_requires_grad.lin1.weight.requires_grad = False
model_requires_grad.lin1.bias.requires_grad = False

# 打印参数的requires_grad属性
print("\n--- 设置 requires_grad = False 后的属性 ---")
print(f"lin0.weight.requires_grad: {model_requires_grad.lin0.weight.requires_grad}")
print(f"lin1.weight.requires_grad: {model_requires_grad.lin1.weight.requires_grad}")
print(f"lin2.weight.requires_grad: {model_requires_grad.lin2.weight.requires_grad}")

# 验证优化器只会更新requires_grad=True的参数
# optimizer = torch.optim.SGD(model_requires_grad.parameters(), lr=0.01)
# 更好的做法是只传入需要更新的参数
# trainable_params = filter(lambda p: p.requires_grad, model_requires_grad.parameters())
# optimizer = torch.optim.SGD(trainable_params, lr=0.01)

行为分析：

通过将model.lin1.weight.requires_grad和model.lin1.bias.requires_grad设置为False，我们成功地冻结了lin1层的参数。在这种情况下，当进行反向传播时，PyTorch会正常计算通过lin2和lin1的梯度，但当遇到lin1的参数时，由于它们的requires_grad为False，其梯度不会被计算和存储。然而，lin1的输出仍然是一个追踪梯度的张量（因为它的输入x来自lin0，且lin0的参数requires_grad为True），因此梯度可以继续反向传播到lin0，使得lin0的参数能够正常更新。

结论： 这是实现精确、选择性层冻结的推荐方法。它允许我们灵活地控制模型中哪些部分的参数参与训练，哪些部分保持不变。

注意事项与最佳实践

1. 全面冻结层参数： 当冻结一个层时，请确保同时设置其所有可学习参数（通常包括weight和bias）的requires_grad为False。
2. 优化器参数： 当你冻结了部分参数后，最好只将requires_grad=True的参数传递给优化器。这可以通过以下方式实现：

   trainable_params = filter(lambda p: p.requires_grad, model.parameters())
   optimizer = torch.optim.Adam(trainable_params, lr=0.001)

   虽然将所有model.parameters()传入优化器也能正常工作（优化器会忽略requires_grad=False的参数），但明确过滤可以提高效率和代码清晰度。

3. model.eval() 的作用： model.eval() 主要用于将模型设置为评估模式，它会禁用Dropout层和BatchNorm层的训练行为（例如，BatchNorm会使用全局均值和方差而不是批次统计量）。它不会自动冻结参数或禁用梯度计算。因此，冻结参数需要单独设置requires_grad = False。
4. torch.no_grad() 的适用场景： torch.no_grad() 仍然是进行推理、验证或在训练循环中不需要梯度计算的特定代码块（例如，计算不参与损失的辅助指标）的理想选择，因为它能节省内存并加速计算。

总结

在PyTorch中冻结模型中间层参数时，理解torch.no_grad()和layer.requires_grad = False之间的区别至关重要。torch.no_grad()是一个上下文管理器，会禁用其作用域内所有计算的梯度追踪，可能意外地影响上游层的梯度流。而直接设置layer.requires_grad = False则是更精确和推荐的方法，它允许我们选择性地冻结特定层的参数，同时保持其他层参数的正常训练。掌握这一技术对于高效地进行模型微调、迁移学习以及各种深度学习实验具有重要意义。

以上就是PyTorch中精确冻结中间层参数的策略与实践的详细内容，更多请关注php中文网其它相关文章！