Pytorch——训练时,冻结网络部分参数的方法

一、原理:

要固定训练网络的哪几层,只需要找到这几层参数(parameter),然后将其 .requires_grad 属性设置为 False 。然后修改优化器,只将不被冻结的层传入。

二、效果

  1. 节省显存:不将不更新的参数传入optimizer
  2. 提升速度:将不更新的参数的requires_grad设置为False,节省了计算这部分参数梯度的时间

三、代码:

.requires_grad 属性设置为 False

python 复制代码
# 根据参数层的 name 来进行冻结
unfreeze_layers = ["text_id"] # 用列表
# 设置冻结参数:
for name, param in model.named_parameters():
    # print(name, param.shape)
    # 错误判定:
    # if name.split(".")[0] in unfreeze_layers: # 不要用in来判定,因为"id"也在"text_id"的in中。
    # 正确判定:
    for unfreeze_layer in unfreeze_layers:
        if name.split(".")[0] != unfreeze_layer:
            param.requires_grad = False
            print(name, param.requires_grad)
        else:
            print(name, param.requires_grad)
python 复制代码
# 冻结整个网络
for param in self.model.parameters():
    param.requires_grad = False
python 复制代码
# 查看冻结参数与否:
for name, param in self.clip_model.named_parameters():
    print(name, param.requires_grad)

修改优化器

python 复制代码
# 只将未被冻结的层,传入优化器
optimizer = optim.SGD(filter(lambda p : p.requires_grad, model.parameters()), lr=1e-2)

四、其他知识

  1. 模型权重冻结:一些情况下,我们可能只需要对模型进行推断,而不需要调整模型的权重。通过调用model.eval(),可以防止在推断过程中更新模型的权重。
  2. with torch.no_grad(): # 禁用梯度计算以加快计算速度。
  3. 训练完train_datasets之后,model要来测试样本了。在model(test_datasets)之前,需要加上model.eval(). 否则的话,有输入数据,即使不训练,它也会改变权值。这是model中含有batch normalization层所带来的的性质。
    eval()时,pytorch会自动把BN和DropOut固定住,不会取平均,而是用训练好的值。不然的话,一旦test的batch_size过小,很容易就会被BN层导致生成图片颜色失真极大。eval()在非训练的时候是需要加的,没有这句代码,一些网络层的值会发生变动,不会固定,你神经网络每一次生成的结果也是不固定的,生成质量可能好也可能不好。
  4. 何时用model.eval() :训练完train样本后,生成的模型model要用来测试样本。在model(test)之前,需要加上model.eval(),否则的话,有输入数据,即使不训练,它也会改变权值。这是model中含有BN层和Dropout所带来的的性质。在eval/test过程中,需要显示地让model调用eval(),此时模型会把BN和Dropout固定住,不会取平均,而是用训练好的值。
  5. 何时用with torch.no_grad():无论是train() 还是eval() 模式,各层的gradient计算和存储都在进行且完全一致,只是在eval模式下不会进行反向传播。而with torch.no_grad()则主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。若想节约算力,可在test阶段带上torch.no_grad()。
  6. with torch.no_grad() 主要是用于停止autograd模块的工作,以起到加速和节省显存的作用。它的作用是将该with语句包裹起来的部分停止梯度的更新,从而节省了GPU算力和显存,但是并不会影响dropout和BN层的行为。如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation/test的结果;而with torch.no_grad()则是更进一步加速和节省gpu空间(因为不用计算和存储梯度),从而可以更快计算,也可以跑更大的batch来测试。

参考文章

  1. 知乎讨论
  2. 第十三章 深度解读预训练与微调迁移,模型冻结与解冻(工具)
  3. 【PyTorch】搞定网络训练中的model.train()和model.eval()模式
相关推荐
蝈蝈tju几秒前
Vibe Coding 正确姿势: 先会指挥, 再让AI干
人工智能·经验分享·ai
想你依然心痛6 分钟前
AI 换脸新纪元:Facefusion 人脸融合实战探索
人工智能·换脸·facefusion·人脸融合
马士兵教育9 分钟前
计算机专业学生入行IT行业,编程语言如何选择?
java·开发语言·c++·人工智能·python
CoderJia程序员甲10 分钟前
GitHub 热榜项目 - 日榜(2026-01-28)
人工智能·ai·大模型·github·ai教程
农场主John19 分钟前
Accelerate_deepspeed使用
pytorch·llm·deepspeed
康谋自动驾驶20 分钟前
高校自动驾驶研究新基建:“实测 - 仿真” 一体化数据采集与验证平台
人工智能·机器学习·自动驾驶·科研·数据采集·时间同步·仿真平台
Ftsom20 分钟前
【6】kilo 上下文管理与压缩机制
人工智能·agent·ai编程·kilo
shangjian00725 分钟前
AI-大语言模型LLM-Transformer架构1-整体介绍
人工智能·语言模型·transformer
机 _ 长30 分钟前
YOLO26 蒸馏改进全攻略:从理论到实战 (Response + Feature + Relation)
人工智能·深度学习·yolo·目标检测·计算机视觉
shangjian00733 分钟前
AI-大语言模型LLM-Transformer架构2-自注意力
人工智能·语言模型·transformer