[python][pytorch]多GPU下的模型保存与加载
说明
在模型训练的时候，往往使用的是多GPU的环境；但是在模型验证或者推理阶段，往往使用单GPU甚至CPU进行运算。那么中间有个保存和加载的过程。下面来总结一下。
多GPU进行训练
首先设置可见的GPU数量，有两种方式可以声明：
- 在shell脚本中声明：
export CUDA_VISIBLE_DEVICES=0,1,2,3
- 在py文件中声明
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
推荐使用前者进行声明，因为后者可能会出现失效的情况。
多GPU模型加载
其次，要将模型分发到不同的GPU。
model = Model(args)
if torch.cuda.is_available() and args.use_gpu:
model= model.cuda()
model = torch.nn.DataParallel(model)
当然这里只是涉及到一个简单的模型并行加载，里面还埋着其他的坑，如果是小数据集且显存够用，完全不用优化，但是如果不够用，我们后面会详细深挖并行中出现坑。
模型保存
等到训练完成之后，需要将模型保存起来。需要注意的是，模型此时保存的是计算图+参数是并行的，但是参数是单GPU的。
state = {
'epoch': epoch,
'model': args.model,
'dataset': args.dataset,
'state_dict': net.module.state_dict() if isinstance(net, nn.DataParallel) else net.state_dict(),
'acc': top1.avg,
'optimizer': optimizer.state_dict(),
}
torch.save(state, filename)
如果服务器环境变化不大，或者和训练时候是同一个GPU环境，直接加载model就不会出现问题，否则建议直接使用参数加载。
模型加载
由于模型训练和部署情况的多样性，大致可以分为以下几种情况：
- 单卡训练，单卡加载部署，单CPU和GPU统一放到这一类。举例：在GPU上训练，在CPU上加载。或者在GPU上训练，在GPU上加载。
这类情况最简单，简单粗暴直接写就行。
model = Model(args)
ckpt = torch.load(args.pretrained_model, map_location='cpu')
state = ckpt['state_dict']
net.load_state_dict(state)
注意map_location的参数，如果在gpu上进行加载，则声明map_location='cuda:0'。如果不声明，可能会报错，input和weight的类型不一致。
- 多卡训练，单卡加载部署。举例：在多GPU上并行训练，在单GPU或CPU上加载。
这种情况要防止参数保存的时候没有加module,那么保存的参数名称是module.conv1.weight，而单卡的参数名称是conv1.weight，这时就会报错，找不到相应的字典的错误。
此时可以通过手动的方式删减掉模型中前几位的名称，然后重新加载。
kwargs={'map_location':lambda storage, loc: storage.cuda(gpu_id)}
def load_GPUS(model,model_path,kwargs):
state_dict = torch.load(model_path,**kwargs)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
return model
-
单卡训练，多卡加载部署。举例：多见于暴发户的情况，一开始只能单卡跑，后来有了多卡，但是单卡的参数有不想浪费。
此时唯有记住一点，因为参数是没有module的，而加载后的参数是有module的，因此需要保证参数加载在模型分发之前。
即保证：
net.load_state_dict(state)在model = torch.nn.DataParallel(model)之前。 -
多卡训练，多卡加载部署。环境如果没有变化，则可以直接加载，如果环境有变化，则可以拆解成第2种情况，然后再分发模型。
