原因:预训练权重层数的键值与新构建的模型中的权重层数名称不吻合,Checkpoint里面的模型是在双卡上训练的,保存的key前面都多一个module.
解决:model = torch.nn.DataParallel(model, device_ids=[0, 1]).cuda()
torch.nn.DataParallel是一种能够将数据分散到多张显卡上从而加快模型训练的方法。它的原理是首先在指定的每张显卡上拷贝一份模型,然后将输入的数据分散到各张显卡上,计算梯度,回传到第一张显卡上,然后再对模型进行参数优化。
注意:不能直接在load_state_dict里面加strict = False来解决此问题,加入strict = False,预训练权重层数的键值与新构建的模型中的权重层数名称不用完全吻合,容易出错。
torch.load_state_dict()函数就是用于将预训练的参数权重加载到新的模型之中,load_state_dict(fsd,strict=False),当strict=True,要求预训练练权重层数的键值与新构建的模型中的权重层数名称完全吻合。
pytorch中state_dict()和load_state_dict()函数配合使用可以实现状态的获取与重载,load()和save()函数配合使用可以实现参数的存储与读取。其中最重要的部分是“字典”的概念,因为参数对象的存储是需要“名称”——“值”对应(即键值对),读取时也是通过键值对读取的。
载入muti-GPU模型:
pretrain_model = torch.load('muti_gpu_model.pth') # 网络+权重# 载入为single-GPU模型gpu_model = pretrained_model.module# 载入为CPU模型model = ModelArch()pretained_dict = pretained_model.module.state_dict()model.load_satte_dict(pretained_dict)
载入muti-GPU权重:
model = ModelArch().cuda() model = torch.nn.DataParallel(model, device_ids=[0]) # 将model转为muti-GPU模式checkpoint = torch.load(model_path, map_location=lambda storage, loc:storage) model.load_state_dict(checkpoint)# 载入为single-GPU模型gpu_model = model.module# 载入为CPU模型model = ModelArch()model.load_state_dict(gpu_model.state_dict())torch.save(cpu_model.state_dict(), 'cpu_model.pth')
载入CPU权重:
# 载入为CPU模型model = ModelArch()checkpoint = torch.load(model_path, map_location=lambda storage, loc:storage) # 载入为single-GPU模型model = ModelArch().cuda() checkpoint = torch.load(model_path, map_location=lambda storage, loc:storage.cuda(0)) model.load_state_dict(checkpoint)# 载入为muti-GPU模型model = ModelArch().cuda() model = torch.nn.DataParallel(model, device_ids=[0, 1]) checkpoint = torch.load(model_path, map_location=lambda storage, loc:storage.cuda(0)) model.module.load_state_dict(checkpoint)
其他报错:
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
解决:把所有tensor都要放在GPU上去