You need to enable JavaScript to run this app.
最新活动
大模型
产品
解决方案
定价
生态与合作
支持与服务
开发者
了解我们

PyTorch中级联神经网络预训练第一网络的冻结方法及eval()有效性验证

解答:不能用eval()冻结权重,正确的冻结方法在这里

首先直接给结论:执行first_model = first_model.eval()并不能冻结第一网络的权重,训练时如果没做其他处理,优化器依然会更新第一网络的参数。

为什么eval()不行?

eval()的作用是把模型切换到评估模式,主要影响的是像BatchNorm、Dropout这类带有训练/评估差异的层:

  • 对于Dropout:评估模式下不会随机丢弃神经元
  • 对于BatchNorm:评估模式下会使用训练阶段统计好的running_mean和running_var,而不是实时计算当前batch的均值方差

eval()完全不会改变参数的可训练状态(也就是requires_grad属性),只要参数的requires_grad=True,优化器就会在反向传播时更新它们。

正确的冻结第一网络的方法

要实现仅训练第二网络,你需要两步操作:

1. 冻结第一网络的所有参数

遍历第一网络的所有参数,将它们的requires_grad设置为False,这样反向传播时就不会计算这些参数的梯度,优化器也不会更新它们:

first_model = get_first_model(load_weights=True)

# 冻结第一网络的所有参数
for param in first_model.parameters():
    param.requires_grad = False

2. 配置优化器只更新第二网络的参数

构建优化器时,只传入第二网络的参数(或者过滤出所有requires_grad=True的参数),确保优化器只关注需要训练的部分:

next_model = next_model_architecture()
network = cascading_model(first_model, next_model)

# 方法1:直接传入第二网络的参数
optimizer = torch.optim.Adam(next_model.parameters(), lr=1e-3)

# 方法2:过滤整个网络中需要训练的参数(适合更复杂的场景)
# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, network.parameters()), lr=1e-3)

(可选但推荐)配合eval()保持预训练模型的状态

如果你的第一网络包含BatchNorm或Dropout层,建议在训练第二网络时把第一网络设为eval()模式,这样这些层会保持预训练时的行为,不会在训练第二网络时改变统计量或随机丢弃:

first_model.eval()

完整修改后的代码示例

# 加载预训练的第一网络并冻结
first_model = get_first_model(load_weights=True)
for param in first_model.parameters():
    param.requires_grad = False
first_model.eval()  # 保持预训练模型的评估模式状态

# 构建第二网络和级联模型
next_model = next_model_architecture()
network = cascading_model(first_model, next_model)

# 优化器仅针对第二网络
optimizer = torch.optim.Adam(next_model.parameters(), lr=1e-3)

# 后续训练循环中,正常执行forward、loss计算、backward、optimizer.step()即可

额外说明

如果之后你想解冻第一网络的部分层(比如微调某几层),只需要把对应层的参数requires_grad重新设为True,并调整优化器的参数列表即可。

内容的提问来源于stack exchange,提问作者John

火山引擎 最新活动