pytorch GPU和CPU模型相互加载方式_F11 - 专业站长和开发者的学习网站

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

GPU训练的模型，在GPU上使用；
CPU训练的模型，在CPU上使用，

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

1	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'