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

Posted 2021-12-12 lishanlu136

在人脸识别任务中，当我定义模型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    # 用完之后清除内存

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

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

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

2. 加载，使用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()