pytorch载入预训练模型后，只想训练个别层怎么办？使用filter(lambda p: p.requires_grad, model.parameters())

有了已经训练好的模型参数，对这个模型的某些层做了改变，如何利用这些训练好的模型参数继续训练

    def init_weights(self, pretrained='',):
        logger.info('=> init weights from normal distribution')
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.001)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
        if os.path.isfile(pretrained):
            pretrained_dict = torch.load(pretrained)
            logger.info('=> loading pretrained model {}'.format(pretrained))
            model_dict = self.state_dict()
            pretrained_dict = {k: v for k, v in pretrained_dict.items()
                               if k in model_dict.keys()}
            #for k, _ in pretrained_dict.items():
            #    logger.info(
            #        '=> loading {} pretrained model {}'.format(k, pretrained))
            model_dict.update(pretrained_dict)
            self.load_state_dict(model_dict)

如果载入的这些参数中，有些参数不要求被更新，即固定不变，不参与训练，需要手动设置这些参数的梯度属性为Fasle，并且在optimizer传参时筛选掉这些参数，更简单的一种方法是在你的优化器中写模型参数的位置改成这样一串代码

optimizer = torch.optim.SGD([{'params':
                              filter(lambda p: p.requires_grad, model.parameters()),
                              'lr': args.LR}],
                              lr=args.LR,
                              momentum=args.MOMENTUM,
                              weight_decay=args.WD,
                              nesterov=args.NESTEROV,
                             )

 

这也是局部finetune的一个方法