python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > pytorch GPU和CPU模型相互加载

pytorch GPU和CPU模型相互加载方式

作者:Overboom

在PyTorch中,保存和加载模型有两种主要方式:直接保存整个模型结构加权重,或者只保存模型的参数,直接保存整个模型的方法简单,但不够灵活,且可能存在模型结构不一致的风险,推荐的做法是只保存模型参数,这种方法需要在加载前定义与原模型结构相同的模型

1 pytorch保存模型的两种方式

1.1 直接保存模型并读取

# 创建你的模型实例对象: model
model = net()
## 保存模型
torch.save(model, 'model_name.pth')

## 读取模型
model = torch.load('model_name.pth')

1.2 只保存模型中的参数并读取

## 保存模型
torch.save({'model': model.state_dict()}, 'model_name.pth')

## 读取模型
model = net()
state_dict = torch.load('model_name.pth')
model.load_state_dict(state_dict['model'])

如何保存模型决定了如何读取模型,一般来选择第二种来保存和读取。

2 GPU / CPU模型相互加载

2.1 单个CPU和单个GPU模型加载

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

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

情况一:CPU -> CPU, GPU -> GPU

这种情况下我们都只用直接用下面的语句即可:

torch.load('model_dict.pth')

情况二:GPU -> CPG/GPU

GPU训练的模型,不知道放在CPU还是GPU运行,两种情况都要考虑

import torch
from torchvision import models

# 加载预训练的GPU模型权重文件
weights_path = 'model_gpu.pth'

# 定义一个与原模型结构相同的新模型
model = models.resnet50()

# 检查是否有可用的CUDA设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将权重映射到相应的设备内存并加载到模型中
weights = torch.load(weights_path, map_location=device)
model.load_state_dict(weights)

# 设置为评估模式
model.eval()

print("Model is successfully loaded and can be used on a", device.type, "!")

情况三:CPU -> CPG/GPU

模型是在CPU上训练的,但不确定要在CPU还是GPU上运行时,两种情况都要考虑

import torch
from torchvision import models

# 加载预训练的CPU模型权重文件
weights_path = 'model_cpu.pth'

# 定义一个与原模型结构相同的新模型
model = models.resnet50()

# 检查是否有可用的CUDA设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 将权重映射到相应的设备内存并加载到模型中
if device.type == 'cuda':
    model.to(device)
    weights = torch.load(weights_path, map_location=device)
else:
    weights = torch.load(weights_path, map_location='cpu')

model.load_state_dict(weights)

# 设置为评估模式
model.eval()

print("Model is successfully loaded and can be used on a", device.type, "!")

总结

以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。

您可能感兴趣的文章:
阅读全文