load_state_dict报错: Error(s) in loading state_dict for XXX: Missing keys(s) in state_dict: ....

1. 报错内容

/opt/tools/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1221 
   1222         if len(error_msgs) > 0:
-> 1223             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
   1224                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1225         return _IncompatibleKeys(missing_keys, unexpected_keys)

RuntimeError: Error(s) in loading state_dict for UNet:
	Missing key(s) in state_dict:................（太多了）
	Unexpected key(s) in state_dict: ................(也很多）

2. 解决方案

train的时候我加了数据并行

# train
model = nn.DataParallel(model)

但是test里没加并行，加上就不会报错了

# test
model = UNet(n_channels=3, n_classes=1)
model = nn.Dat