pytorch冻结参数训练的坑

由于项目需要训练一个主干网络接多个分支的模型,所以先训练一个主干网络加第一个分支,再用另外的数据训练第二个分支,训练的过程中需要冻结主干网络部分,后面的分支训练过程也一样需要冻结主干网络部分。

冻结模型的方式

python 复制代码
for name, para in model.named_parameters():
      # 冻结backbone的权重
	if name.split(".")[0] == "backbone":
          para.requires_grad = False        # 或者用para.requires_grad_(False),一个是通过属性直接赋值,一个是通过函数赋值
    else:
          para.requires_grad = True
python 复制代码
# 可以打印需要更新梯度的参数
for name, value in model.named_parameters():
    print(name, "\t更新梯度:",value.requires_grad)

坑1:这样做并不能冻结batchnorm层的参数,所以需要在训练中手动冻结。如:

python 复制代码
def fix_bn(m):
    classname = m.__class__.__name__
    if classname.find('SyncBatchNorm') != -1 or classname.find('InstanceNorm2d') != -1 or classname.find('BatchNorm2d') != -1:          #SyncBatchNorm, InstanceNorm2d
        if m.num_features in [32, 64, 96, 128, 256, 384, 768, 192, 1152, 224]:      # 需要冻结的BN层的通道数
            m.eval()

def train():
	for epoch in range(max_epoch):
		model.train()
		if args.freeze:
			model.apply(fix_bn)
			model.backbone[5][0].block[0][1].eval()   # 假如需要冻结的BN层通道数和不需要冻结的BN层通道数一样,则需要单独写
		for batch_idx, (data, target) in enumerate(train_loader):
			...
			

坑2:用了冻结训练(freeze)就不要用EMA方式更新模型了,不然收敛缓慢不说,还会造成前面冻结的参数产生变化,可以从EMA的代码看出端倪:

python 复制代码
class EMA:
    def __init__(self, model, decay=0.9999):
        super().__init__()
        import copy
        self.decay = decay
        self.model = copy.deepcopy(model)

        self.model.eval()

    def update_fn(self, model, fn):
        with torch.no_grad():
            e_std = self.model.state_dict().values()
            #m_std = model.module.state_dict().values()   # multi-gpu
            m_std = model.state_dict().values()          # single-gpu
            for e, m in zip(e_std, m_std):
                e.copy_(fn(e, m))

    def update(self, model):
        self.update_fn(model, fn=lambda e, m: self.decay * e + (1. - self.decay) * m)

可以看出EMA的方式更新模型方式,大部分是结合上一个模型的参数的,即:

复制代码
model_update = decay*model(t-1) + (1-decay)*model(t)       # model(t-1) 代表上一次迭代模型的参数,model(t)代表当前迭代得到的模型参数

虽然冻结了backbone的参数,阻止了梯度在backbone中反向传播,但参数由于经过如上乘法及加法运算,由于精度原因,还是会发生微小变化,虽然训练次数增加,这个变化会扩大,从而达不到冻结训练的效果。而且从计算公式可以看出来,采用EMA的方式更新模型参数,参数会更新得很慢,会造成网络难以学习的"错觉"。我在这里困住了3天,有怀疑过是否是网络设计问题,是否是多GPU同步的问题,是否是参数设置,如学习率过小,权重衰减过大,或者dropout设置过大等等,最终一步一步排除定位到EMA的问题。

以这次的经验来看,EMA只适合在上一次训练得到模型的基础上,这一次加了额外的数据,需要在上一次的基础上做微调的情况。

相关推荐
运器1233 分钟前
【一起来学AI大模型】PyTorch DataLoader 实战指南
大数据·人工智能·pytorch·python·深度学习·ai·ai编程
音元系统6 分钟前
Copilot 在 VS Code 中的免费替代方案
python·github·copilot
超龄超能程序猿18 分钟前
(5)机器学习小白入门 YOLOv:数据需求与图像不足应对策略
人工智能·python·机器学习·numpy·pandas·scipy
卷福同学18 分钟前
【AI编程】AI+高德MCP不到10分钟搞定上海三日游
人工智能·算法·程序员
帅次26 分钟前
系统分析师-计算机系统-输入输出系统
人工智能·分布式·深度学习·神经网络·架构·系统架构·硬件架构
AndrewHZ1 小时前
【图像处理基石】如何入门大规模三维重建?
人工智能·深度学习·大模型·llm·三维重建·立体视觉·大规模三维重建
5G行业应用1 小时前
【赠书福利,回馈公号读者】《智慧城市与智能网联汽车,融合创新发展之路》
人工智能·汽车·智慧城市
悟空胆好小1 小时前
分音塔科技(BABEL Technology) 的公司背景、股权构成、产品类型及技术能力的全方位解读
网络·人工智能·科技·嵌入式硬件
探讨探讨AGV1 小时前
以科技赋能未来,科聪持续支持青年创新实践 —— 第七届“科聪杯”浙江省大学生智能机器人创意竞赛圆满落幕
人工智能·科技·机器人
cwn_1 小时前
回归(多项式回归)
人工智能·机器学习·数据挖掘·回归