python多显卡支持_[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种情况,然后再分发模型。

你可能感兴趣的:(python多显卡支持)