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()模式
相关推荐
边缘计算社区1 小时前
FPGA与边缘AI:计算革命的前沿力量
人工智能·fpga开发
飞哥数智坊1 小时前
打工人周末充电:15条AI资讯助你领先一小步
人工智能
Tech Synapse1 小时前
基于CARLA与PyTorch的自动驾驶仿真系统全栈开发指南
人工智能·opencv·sqlite
layneyao1 小时前
深度强化学习(DRL)实战:从AlphaGo到自动驾驶
人工智能·机器学习·自动驾驶
海特伟业2 小时前
隧道调频广播覆盖的实现路径:隧道无线广播技术赋能行车安全升级,隧道汽车广播收音系统助力隧道安全管理升级
人工智能
CareyWYR2 小时前
每周AI论文速递(250421-250425)
人工智能
追逐☞2 小时前
机器学习(10)——神经网络
人工智能·神经网络·机器学习
winner88812 小时前
对抗学习:机器学习里的 “零和博弈”,如何实现 “双赢”?
人工智能·机器学习·gan·对抗学习
Elastic 中国社区官方博客2 小时前
使用 LangGraph 和 Elasticsearch 构建强大的 RAG 工作流
大数据·数据库·人工智能·elasticsearch·搜索引擎·ai·全文检索
娃娃略2 小时前
【AI模型学习】双流网络——更强大的网络设计
网络·人工智能·pytorch·python·神经网络·学习