在迁移学习中，要使用预训练权重。在训练好的网络权重中，使用的是1000类的分类任务。而我们自己的分类任务中没有1000类，需要修改网络。如果将MobileNetV2(num_classes=5)，我们的网络在最后输出是5类，与1训练好的1000类不符。加载时就会报错。这时候就要将训练参数中的最后输出1000类的那层删除掉。

    # create modelnet = MobileNetV2(num_classes=5)# load pretrain weights# download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pthmodel_weight_path = "./mobilenet_v2.pth"assert os.path.exists(model_weight_path), "file {} dose not exist.".format(model_weight_path)pre_weights = torch.load(model_weight_path, map_location='cpu')# delete classifier weightspre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)# freeze features weightsfor param in net.features.parameters():param.requires_grad = False

然后将参数冻结，只训练最后一层参数。

参考视频：如何在Pytorch中载入部分权重_哔哩哔哩_bilibili