U-Net网络使用预训练好的权重时遇到的问题与解决办法

1、自己电脑有gpu，但是训练时出现以下问题

raise RuntimeError('Attempting to deserialize object on a CUDA '
RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_avwith map_location=torch.device('cpu') to map your storages to the CPU.

解决办法：将unet.py里面的True改为False，并且改动以下代码（加上map_location=“cpu”）

 state_dict = torch.load(self.model_path,map_location="cpu")

2、

RuntimeError: Error(s) in loading state_dict for Unet:
        size mismatch for final.weight: copying a param with shape torch.Size([21, 64, 1, 1]) from checkpoint, the shape in current model is torch.Size([2, 64, 1, 1]). 
        size mismatch for final.bias: copying a param with shape torch.Size([21]) from checkpoint, the shape in current model is torch.Size([2]).

即加载的模型参数和初始化的Unet结构不匹配的问题。
解决办法：
我使用的是voc数据集的预训练好的权重，所以num_classes应该是21，刚开始我写的是2，改完unet.py的数值后运行成功可以预测

     "model_path"        : 'model_data/unet_voc.pth',
        "model_image_size"  : (512, 512, 3),
        "num_classes"       : 21,
        "cuda"              : False,