def load_model(self):
if self.args.snapshot:
print('load model...')
self.load_state_dict(torch.load(self.args.snapshot), strict=False)
'''
if self.args.snapshot:
print('load model for localization...')
new_state_dict={}
for k, v in torch.load(self.args.snapshot,map_location=torch.device('cuda:0')).items():
new_state_dict[k[7:]] = v #键值包含‘module.’ 则删除
self.load_state_dict(new_state_dict,strict=False)
'''
多卡训练时权重名字会带有module,有时候会加载不进去,尤其是strict=false的情况下。
注释中适用多卡加载模型的情况