pytorch多卡训练和单卡训练加载预训练模型时可能出现的问题

最新推荐文章于 2023-09-25 18:09:20 发布

WSOL_fym

最新推荐文章于 2023-09-25 18:09:20 发布

阅读量472

点赞数

文章标签： pytorch 深度学习 python

本文链接：https://blog.csdn.net/m0_57522162/article/details/129855788

版权

    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的情况下。

注释中适用多卡加载模型的情况