pytorch优化器传入两个模型的参数/保存和加载两个模型的参数

最新推荐文章于 2024-08-06 17:12:31 发布

lishanlu136

最新推荐文章于 2024-08-06 17:12:31 发布

阅读量6.6k

点赞数 7

分类专栏： # Pytorch 文章标签： pytorch 深度学习计算机视觉

本文链接：https://blog.csdn.net/lishanlu136/article/details/121284421

版权

Pytorch 专栏收录该内容

6 篇文章 0 订阅

订阅专栏

该博客探讨了在人脸识别任务中如何配置Arcface_loss，并详细说明了如何根据权重更新规则将模型参数分为多个组，以便在训练过程中对它们进行独立管理和更新。博客内容涉及了如何初始化和管理两个模型的优化器，包括对权重衰减的处理，以及如何保存和加载模型参数。此外，还介绍了如何根据参数类型和需求将模型参数分组，确保不同部分的参数得到适当的训练和调整。

摘要由CSDN通过智能技术生成

在人脸识别任务中，当我定义模型backbone后，用到Arcface_loss，但这个Arcface_loss也是用nn.Module模块写的，所以实例化出来也是一个网络，而且原论文中，Arcface_loss还是以backbone权重参数10倍的权重衰减方式更新，需要单独以不同的方式训练，且网络中bn层也是不需要权重衰减的。由于这些原因，我们就需要把网络参数分开。

分别传入两个模型的参数

如果我们只有一个模型model，我们一般就是这样初始化优化器的：

optimizer = torch.optim.SGD(model.parameters(), lr, momentum=0.9)

现在有两个模型model和head了，要把参数都加进优化器，并且以不同lr来训练参数，我们就可以这样来初始化优化器：

optimizer = torch.optim.SGD([{'params': model.parameters()},
                             {'params': head.parameters(), 'lr': 1e-4}],
                             lr=0.01, momentum=0.9)

还可以将整个模型按参数需不要训练更新，权重衰减来分组：

def separate_paras(modules):
    parses_bn, parses_w, parses_bias = [], [], []
    for k, v in modules.named_modules():
        if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
            parses_bias.append(v.bias)  # biases
        if isinstance(v, nn.BatchNorm2d):
            parses_bn.append(v.weight)  # no decay
        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
            parses_w.append(v.weight)  # apply decay

    return parses_w, parses_bias, parses_bn

paras_w, paras_bias, paras_bn = separate_paras(model)
optimizer = torch.optim.SGD(paras_bn, lr=args.lr_model, momentum=args.momentum, nesterov=True)  # 只需设置学习率
optimizer.add_param_group({'params': paras_w[:-1], 'weight_decay': args.weight_decay})  # backbone权重参数，需要权重衰减
optimizer.add_param_group({'params': [paras_w[-1]]+[head.kernel], 'weight_decay': args.weight_decay*1.5}) #特殊层，需要单独定义权重衰减参数
optimizer.add_param_group({'params': paras_bias})  # 偏置
del paras_w, paras_bias, paras_bn    # 用完之后清除内存

保存/加载两个模型的参数

保存，使用torch.save(state, filename)

filename = './facemodel.pth'
state = {'model':model.state_dict(),
         'head':head.state_dict()}
torch.save(state, filename)

加载，使用torch.load()，并使用字典进行匹配加载

load_name = './facemodel.pth'
checkpoint = torch.load(load_name)
model.load_state_dict(checkpoint['model'])
head.load_state_dict(checkpoint['head'])
model.cuda()
head.cuda()