在单机多GPU上加载在单个GPU上训练的预训练参数

Question

我尝试在一台有多个GPU的机器上加载一个只用单个GPU训练的预训练参数，但出现了“缺少键”和“意外键”等错误。

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)

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())))

Answer 1

“缺失键错误”意味着你的模型结构在尝试寻找一些在你加载的文件中不存在的权重。“意外键”则表示你加载的文件里有一些模型不需要的键。这个问题很可能是因为你的模型结构在保存权重后发生了变化，导致一些状态字典的键被重命名了。