PyTorch在CentOS上的模型保存与加载方法

import torch
import torch.nn as nn

# 定义模型架构
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# 实例化模型
model = SimpleModel()

# 假设模型已完成训练
# 保存模型到文件 'model.pth'
torch.save(model.state_dict(), 'model.pth') # 保存模型参数

# 加载模型参数
model = SimpleModel() # 重新创建模型实例
model.load_state_dict(torch.load('model.pth'))
model.eval() # 设置模型为评估模式

# 将模型转移到合适的设备 (GPU 或 CPU)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 进行预测
input_data = torch.randn(1, 10).to(device) # 示例输入数据，需与设备保持一致
output = model(input_data)

import torch
import torch.nn as nn

# 模型定义
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = SimpleModel()

# 模拟训练过程 (此处省略)

# 保存模型参数
torch.save(model.state_dict(), 'model.pth')

# 加载模型参数
model = SimpleModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 选择设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 进行预测
input_data = torch.randn(1, 10).to(device)
output = model(input_data)
print(output)