pytorch 单卡加载多卡模型module

 

'''
实际上就是字典的操作，那么字典的操作，哪些层不要，打印出层的名字就可以了，例如
模型参数的某些层的权重不要，那么重构一个字典参数就可以了，for k in torch.load("**.pth').keys():   打印出来按照名字删除 键值对，
所以字典的pop删除操作也是可以的， 多卡训练参数多了一个module

'''

# pretrained_dict = torch.load("checkpoint.pth")['state_dict']  #如果存储的是 多个字典信息


pretrain = torch.load("/home/imagenet.pth")
new_state_dict = {}#OrderedDict()
for k,v in pretrain.items():
    # if "classifier" in k:  #最后分类层的参数是classeifer ,不需要这个模型参数
        # continue
    new_state_dict[k[7:]] = v  #remove `module.`  #模型k 有module 不要
model.load_state_dict(new_state_dict, strict=False)  #strict =False ,模型参数和模型不一致可以加载

 

checkpoint = torch.load('checkpoint_backend.pth.tar' )

model = EfficientNet.from_name('efficientnet-b4', num_classes=472)
model = torch.nn.DataParallel(model).cuda()
model.load_state_dict(checkpoint["state_dict"])

model.cuda()

model.eval()