pytorch GPU和CPU模型相互加载

2023-07-21 16:19| 来源: 网络整理| 查看: 265

一模型保存加载方法

注意：torch.load 进行加载模型的时候，会根据模型训练方式（CPU训练或者GPU训练），自动选择加载设备

1.1、直接保存模型（参数+网络） # 保存模型 torch.save(model, 'model.pth') # 加载模型 model = torch.load('model.pth') 1.2、只保存参数；（官方推荐）

由于保存整个模型将耗费大量的存储，故官方推荐只保存参数，然后在建好模型的基础上加载

# 保存模型参数 torch.save(model.state_dict(), path) # 加载模型参数 model.load_state_dict(torch.load(path))

特别地，如果还想保存某一次训练采用的优化器、epochs等信息，可将这些信息组合起来构成一个字典，然后将字典保存起来：

# 要保存的模型参数、optimizer、epoch 的字典state state = {'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch} # **************** 模型保存**************** # 保存字典 state torch.save(state, path) # **************** 模型加载**************** # 加载保存的字典 state checkpoint = torch.load(path) # 依次加载模型参数、optimizer、epoch model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) epoch = checkpoint(['epoch'])

需要注意的是，只保存参数的方法在加载的时候要事先定义好跟原模型一致的模型，并在该模型的实例对象(假设名为model)上进行加载，即在使用上述加载语句前已经有定义了一个和原模型一样的Net, 并且进行了实例化 model=Net( ) 。

1.3、model.state_dict() 和 model.load_state_dict() 函数解释

model.state_dict()， model.load_state_dict()、model.parameters() 的解释请看：

PyTorch模型保存深入理解 - 简书

关于model.parameters(), model.named_parameters(), model.children(), model.named_children(), model.modules(), model.named_modules() 请看：

PyTorch中的model.modules(), model.children(), model.named_children(), model.parameters(), model.nam... - 简书

二 GPU和CPU模型相互加载 1、单CPU 和单GPU 模型相互加载

pytorch允许把在GPU上训练的模型加载到CPU上，也允许把在CPU上训练的模型加载到GPU上。

加载模型参数的时候，在GPU和CPU训练的模型是不一样的，这两种模型是不能混为一谈的，下面分情况进行操作说明。

情况一：模型参数(CPU->CPU,GPU->GPU)

模型是GPU模型，预加载的训练参数也是GPU；模型是CPU模型，预加载的训练参数也是CPU，这种情况下我们都只用直接用下面的语句即可：

torch.load('model_dict.pth') 情况二：模型参数(GPU->CPU)

模型是CPU模型，预加载的训练参数却是GPU，那么需要将先将GPU模型参数加载到 CPU上：

torch.load('model_dict.pkl', map_location=lambda storage, loc: storage) torch.load('model_dict.pkl', map_location='cpu') # pytorch0.4.0及以上版本情况三：模型参数(CPU->GPU, CPU->GPU1)

模型是GPU或者GPU1模型，预加载的训练参数却是CPU，那么需要将CPU模型参数加载到 GPU上：

# CPU->GPU torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda) # CPU->GPU1 模型是GPU1，预加载的训练参数却是CPU： torch.load('model_dic.pkl', map_location=lambda storage, loc: storage.cuda(1)) 2、多GPU上训练的模型保存和加载 2.1 多GPU保存模型保存

多GPU模式下的情况：

# 以下两种模型都是在多GPU上训练, 只有当（GPU >=2）时候才使用 model= torch.nn.DaraParallel(model) # 数据并行 model = torch.nn.parallel.DistributedDataParallel(model) #分布式训练

多GPU训练模型保存参数时，可以去掉 “module.”前缀，这样以后CPU或者单GPU模式加载模型时，相对比较方便。

a、多GPU保存参数与网络： torch.save(model.module,'model.pth') b、多GPU只保存参数：

model.module 操作可以去掉多GPU训练模型时候加的module前缀

#方式1----只保存参数， model.module操作可以去掉多GPU保存参数时候加的module前缀 torch.save(model.module.state_dict(), 'model.pth') 2.2 加载多GPU模型参数（去掉“module.”前缀）

在单gpu或cpu上加载一个在多gpu上训练的模型时，会报类似如下错误：

Unexpected key(s) in state_dict: "module.pretrained.conv1.0.weight", "module.pretrained.conv1.1.weight"

这是因为在多GPU上训练的模型保存时，在参数名前多加了一个“module.”前缀，保存时并没有把“module”前缀去掉，因此加载参数的时候把这个前缀去掉就行了：

from collections import OrderedDict # 模型初始化， CPU上模型 model = net() # 模型参数加载到的设备 device = torch.device('cpu') # device = torch.device('cuda:0') # 多GPU 模型参数加载到 CPU上，这里xxx.pth保存的只是参数 state_dict = torch.load('xxx.pth', map_location=device) # 多GPU 模型参数去掉 module state_dict_new = OrderedDict() for k, v in state_dict.items(): name = k[7:] # 去掉 `module.` # name = k.replace(“module.", "") state_dict_new[name] = v # 模型加载参数（去掉module） model.load_state_dict(state_dict_new) 2.3 多GPU模型参数直接加载

2.2 节主要是在保存和加载多GPU模型时，去掉“module.”前缀，这一节主要讲，另一种加载多GPU模型的方法

Case-1. 载入为多GPU网络+权重，模型为单GPU或者CPU

load() 加载多GPU网络+权重后，获取module部分权重moduel_part，然模型就可以直接加载moduel_part了，此时模型CPU或者GPU都可以

# 模型初始化 model = Mode() # CPU 模型 model = Mode().cuda() # 单GPU 模型 # 加载多GPU 网络+权重 pretrained_model = torch.load('muti_gpus_model.pth') # 只获取 module 部分的数据，可以用于单GPU或者CPU模型加载 moduel_part = pretrained_model.module # 获取module部分的参数字典，可以用于保存或者加载模型参数 pretained_dict = pretained_model.module.state_dict() # 加载模型参数 model.load_state_dict(pretained_dict) # cpu模型存储， model.state_dict() 是只保存模型权重 torch.save(model.state_dict(), 'mode.pth') Case-2.载入多GPU权重，模型为多GPU模型

多GPU模型模型初始化后，多GPU权重可以直接加载

model = Model().cuda(0) # 网络结构 # 将model转为muit-gpus模式 model = torch.nn.DataParallel(model, device_ids=[0]) # 载入多GPU 网络 weights checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) # 用多GPU weights 直接初始化网络多GPU模型 model.load_state_dict(checkpoint) Case-3.载入CPU权重

注意：第三个多GPU模型 weight参数载入使用的是 model.module.load_state_dict(checkpoint)

（1） # 模型放到 CPU model = Model() # cpu 模型直接载入 weights checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage) ****************************************** （2） # 模型放到单GPU model = Model().cuda() # 载入weights 到 cuda 中 checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage.cuda(0)) # 单GPU 模型直接载入 cuda的weights model.load_state_dict(checkpoint) # 用weights初始化网络 ****************************************** （3） model = Model().cuda() # 网络结构 # 模型放到 muti-gpus 中 # device_ids根据自己需求改！ model = torch.nn.DataParallel(model, device_ids=[0, 1]) # 载入weights 到 cuda 中 checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage.cuda(0)) # 多GPU模型 model.module 模块直接载入 cuda的weights model.module.load_state_dict(checkpoint) # 用weights初始化网络

【本文地址】

公司简介

联系我们