CentOS PyTorch如何进行模型保存与加载

在CentOS上使用PyTorch进行模型的保存与加载是一个常见的操作。以下是详细的步骤和示例代码，帮助你完成这一任务。

安装PyTorch

首先，确保你已经安装了PyTorch。你可以参考PyTorch官方网站的安装指南来安装适合你系统的版本。

保存模型

在PyTorch中，你可以使用torch.save()函数来保存整个模型或模型的参数。

保存整个模型

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)

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

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

# 保存整个模型
torch.save(model, 'model.pth')

保存模型参数

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

加载模型

加载模型时，你需要根据保存的方式来选择加载整个模型还是仅加载模型参数。

加载整个模型

# 加载整个模型
loaded_model = torch.load('model.pth')

# 使用模型进行预测
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

加载模型参数

# 创建一个新的模型实例
new_model = SimpleModel()

# 加载模型参数
new_model.load_state_dict(torch.load('model_state_dict.pth'))

# 使用模型进行预测
input_data = torch.randn(1, 10)
output = new_model(input_data)
print(output)

注意事项

1. 模型类定义：在加载模型参数时，必须确保模型类的定义与保存模型时的定义完全一致。
2. 设备兼容性：如果模型是在GPU上训练的，而你在CPU上加载，可能会遇到设备不匹配的问题。可以使用map_location参数来解决这个问题。

   # 如果模型是在GPU上训练的，而你在CPU上加载
   model = torch.load('model.pth', map_location=torch.device('cpu'))
   

示例代码总结

以下是一个完整的示例代码，展示了如何在CentOS上保存和加载PyTorch模型：

import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 5)

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

# 创建模型实例并进行训练（这里省略训练过程）
model = SimpleModel()

# 保存整个模型
torch.save(model, 'model.pth')

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

# 加载整个模型
loaded_model = torch.load('model.pth')

# 使用模型进行预测
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

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

# 使用模型进行预测
input_data = torch.randn(1, 10)
output = new_model(input_data)
print(output)

通过以上步骤，你可以在CentOS上轻松地保存和加载PyTorch模型。