[python][pytorch]多GPU下的模型保存与加载

说明

在模型训练的时候，往往使用的是多GPU的环境；但是在模型验证或者推理阶段，往往使用单GPU甚至CPU进行运算。那么中间有个保存和加载的过程。下面来总结一下。

多GPU进行训练

首先设置可见的GPU数量，有两种方式可以声明：

1. 在shell脚本中声明：

export CUDA_VISIBLE_DEVICES=0,1,2,3

2. 在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就不会出现问题，否则建议直接使用参数加载。

模型加载

由于模型训练和部署情况的多样性，大致可以分为以下几种情况：

1. 单卡训练，单卡加载部署，单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的类型不一致。

2. 多卡训练，单卡加载部署。举例：在多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

3. 单卡训练，多卡加载部署。举例：多见于暴发户的情况，一开始只能单卡跑，后来有了多卡，但是单卡的参数有不想浪费。
   此时唯有记住一点，因为参数是没有module的，而加载后的参数是有module的，因此需要保证参数加载在模型分发之前。
   即保证：
   net.load_state_dict(state)在model = torch.nn.DataParallel(model)之前。

4. 多卡训练，多卡加载部署。环境如果没有变化，则可以直接加载，如果环境有变化，则可以拆解成第2种情况，然后再分发模型。