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

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

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

## 保存模型
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'])

torch.load('model_dict.pth')

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, "!")

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, "!")