PyTorch中冻结中间层参数的深度解析与实践-Python教程-PHP中文网

PyTorch中冻结中间层参数的深度解析与实践

霞舞

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

原创

547人浏览过

PyTorch中冻结中间层参数的深度解析与实践

本教程深入探讨了在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模型中各层的训练状态，从而实现更复杂的训练策略，例如微调预训练模型或进行部分模型的更新。

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