pytorch GPU和CPU模型相互加载

一 模型保存加载方法

注意：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初始化网络