pytorch加载预训练模型并固定部分层，训练一段时间再取消固定从头训练

首先加载预训练模型，将不需要训练的层requires_grad设置为False.(如果想要只训练分类层，记得分类成换成与之前不同的名字)

state_dict_best = torch.load('./xxx.pth')
model.load_state_dict(state_dict_best ,strict = False)  # strict=False可以忽略加载预训练模型和自己不同的层。

for k,v in model.named_parameters():

    if k!='xxx.weight' and k!='xxx.bias' :

        v.requires_grad=False  # 不计算梯度实现了参数固定

optimizer = make_optimizer(model = model,...)

scheduler = WarmupMultiStepLR(optimizer,...)

需要注意训练一段时间如果想再打开之前冻结的层，只要model的reauire_grade设置为True.
同时不要忘记将优化器再重新加载一遍，否则虽然设置为True,依然还是固定训练。代码如下。

for k,v in model.named_parameters():
        v.requires_grad=True  # 固定层打开
optimizer = make_optimizer(model = model,...)

scheduler = WarmupMultiStepLR(optimizer,...)

if use_gpu:
  model = nn.Dataparallel(model).cuda()