Pytorch冻结部分模型参数

在对模型进行训练时，有时会需要冻结模型的一部分参数，不对其进行更新，而只更新剩余部分的参数。

tensor.requires_grad属性

在pytorch0.4版本之后，对Variable和tensor进行了合并，通过设定tensor.requires_grad属性可以控制是否计算tensor的梯度。

x = torch.ones(1)  # create a tensor with requires_grad=False (default) x.requires_grad
x.requires_grad

默认情况下，创建的tensor的requires_grad属性为False，而创建的module的parameters的requires_grad为True。

Optimizer控制参数的更新

在构建优化器实例时，需要传入需要进行更新的参数：

optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr = 0.0001)

因而，可以只向optimizer中传入不需要冻结的参数。

具体做法

将需要冻结的参数的requires_grad属性设置为False。

model_ft = models.resnet50(pretrained=True)  #读入resnet50模型
ct = 0
for child in model_ft.children():
ct += 1
if ct < 7:
    for param in child.parameters():
        param.requires_grad = False

从optimizer更新的参数中剔除requires_grad属性为False的参数。

optimizer.SGD(filter(lambda p: p.requires_grad, model_ft.parameters()), lr=1e-3)

值得注意的是，这里的optimizer声明需要放在param.requires_grad = False后面。

实战

在这个例子中，我们要考虑一种情况，就是我需要先冻结模型，训练几个epoch之后，进行解冻。另外，因为冻结训练的epoch通常比较少，所以不考虑resume的情况，这里只考虑训练到解冻模型时的resume。

首先，冻结模型参数，并声明优化器。

 for param in self.unet.module.backbone.parameters():
        param.requires_grad = False
self.optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.unet.module.parameters()), self.lr, [self.beta1, self.beta2])

训练特定epoch之后，解冻模型参数，将参数添加到优化器。

for param in self.unet.module.backbone.parameters():
       param.requires_grad = True
self.optimizer.add_param_group({'params':self.unet.module.backbone.parameters()})

再训练若干epoch之后，保存模型：

state = {'epoch': epoch,
     'state_dict': self.unet.module.state_dict(),
     'max_dice': self.max_dice,
     'optimizer' : self.optimizer.state_dict(),
     'lr' : self.lr}
torch.save(state, 'checkpoint.pth')

若要resume(重新加载模型参数训练)，这里只考虑训练到训练到解冻模型的resume。如下：

def load_checkpoint(self):
    # Load the pretrained Encoder
    weight_path = os.path.join(self.save_path, self.resume)
    if os.path.isfile(weight_path):
        checkpoint = torch.load(weight_path)
        # 加载模型的参数，学习率，优化器，开始的epoch，最小误差等
        if torch.cuda.is_available:
            self.unet.module.load_state_dict(checkpoint['state_dict'])
        else:
            self.unet.load_state_dict(checkpoint['state_dict'])
        self.lr = checkpoint['lr']
        self.start_epoch = checkpoint['epoch']
        self.max_dice = checkpoint['max_dice']
        self.optimizer.load_state_dict(checkpoint['optimizer'])

        print('%s is Successfully Loaded from %s' % (self.model_type, weight_path))
        write_txt(self.save_path, '%s is Successfully Loaded from %s' % (self.model_type, weight_path))
    else:
        raise FileNotFoundError("Can not find weight file in {}".format(weight_path))

for param in self.unet.module.backbone.parameters():
    param.requires_grad = False
self.optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.unet.module.parameters()), self.lr, [self.beta1, self.beta2])
if self.resume:
    for param in self.unet.module.backbone.parameters():
        param.requires_grad = True
    self.optimizer.add_param_group({'params': self.unet.module.backbone.parameters()})
    self.load_checkpoint()