PyTorch中选择性冻结神经网络中间层参数的策略与实践-Python教程-PHP中文网

PyTorch中选择性冻结神经网络中间层参数的策略与实践

花韻仙語

发布： 2025-08-22 15:48:01

原创

299人浏览过

PyTorch中选择性冻结神经网络中间层参数的策略与实践

本文深入探讨了在PyTorch神经网络中选择性冻结中间层参数的两种主要方法：使用torch.no_grad()上下文管理器和直接设置参数的requires_grad属性为False。通过详细分析和代码示例，文章阐明了这两种方法在梯度计算和参数更新方面的不同影响，强调了requires_grad = False是实现精确层冻结的推荐方案，以确保只有目标层被冻结而其他层能够正常训练。

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：
- 优点：实现精确的层冻结，只冻结目标层的参数，而允许梯度流通过并更新其前置层。这是在迁移学习或需要精细控制训练过程时的推荐方法。
- 注意事项：
  1. 确保在定义优化器时，只将requires_grad = True的参数传递给优化器，例如：optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01)。否则，优化器可能会尝试更新那些不计算梯度的参数，导致不必要的计算或潜在错误。
  2. 冻结操作通常在模型实例化后、训练循环开始前进行。
  3. 可以通过打印param.requires_grad来验证参数的梯度计算状态，并通过比较训练前后的param.data来验证参数是否实际更新。