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


    说明

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

    多GPU进行训练

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

    1. 在shell脚本中声明:
    export CUDA_VISIBLE_DEVICES=0,1,2,3
    
    1. 在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的类型不一致。

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

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

  • 相关阅读:
    [C]recursion递归计算阶乘
    [Python]reduce function & lambda function & factorial
    [C/JAVA] ceil, floor
    OC项目调用C++
    Xcode 代码注释
    百度云加速器
    UITableView和MJReFresh结合使用问题记录
    OC 类的load方法
    JLRoutes笔记
    推送通知项目记录
  • 原文地址:https://www.cnblogs.com/wildkid1024/p/13025352.html
Copyright © 2020-2023  润新知