单机多GPU加载单GPU预训练参数遇Missing keys等错误求助

在单机多GPU环境下加载单GPU训练得到的预训练参数时，出现了Missing keys和Unexpected keys错误。

相关代码 #

1. Backbone初始化代码 #

backbone_cfg = dict(
            embed_dim=embed_dim,
            depths=depths,  
            num_heads=num_heads, 
            window_size=window_size,
            ape=False,
            drop_path_rate=0.3,
            patch_norm=True,
            use_checkpoint=False,
            frozen_stages=frozen_stages
        )

self.backbone = SwinTransformer(**backbone_cfg)

2. 最初的多GPU权重加载函数 #

def init_weights_multiGPUs(self, pretrained=None):
             if pretrained is not None:
                if dist.get_rank() == 0:
                    self.backbone.load_state_dict(torch.load(pretrained))
                dist.barrier()

其中pretrained为预训练参数路径。

已尝试的解决方法（未生效） #

修改后的权重加载函数：

def init_weights_multiGPUs(self, pretrained = None) :
        print(f'== Load encoder backbone on multiGPUs from: {pretrained}')
        if isinstance(self.backbone, torch.nn.parallel.DistributedDataParallel):
            self.backbone = self.backbone.module
        self.backbone.load_state_dict(torch.load(pretrained, map_location='cuda:{}'.format(torch.cuda.current_device())))

内容的提问来源于stack exchange，提问作者Mingshuai Zhao