Ubuntu PyTorch如何保存和加载模型

在Ubuntu上使用PyTorch保存和加载模型的步骤如下：

保存模型

1. 定义模型： 首先，你需要定义你的模型。例如，一个简单的线性回归模型：

   import torch
   import torch.nn as nn
   
   class LinearRegressionModel(nn.Module):
       def __init__(self, input_dim, output_dim):
           super(LinearRegressionModel, self).__init__()
           self.linear = nn.Linear(input_dim, output_dim)
   
       def forward(self, x):
           return self.linear(x)
   

2. 实例化模型： 创建模型的实例，并将其移动到GPU（如果可用）：

   model = LinearRegressionModel(input_dim=10, output_dim=1)
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
   

3. 训练模型： 训练你的模型（这里省略了训练代码）。

4. 保存模型： 使用torch.save()函数保存模型的状态字典：

   torch.save(model.state_dict(), 'model.pth')
   

加载模型

1. 定义相同的模型结构： 加载模型之前，你需要定义与保存时相同的模型结构：

   model = LinearRegressionModel(input_dim=10, output_dim=1)
   device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
   model.to(device)
   

2. 加载模型状态字典： 使用torch.load()函数加载模型的状态字典，并将其加载到模型中：

   model.load_state_dict(torch.load('model.pth', map_location=device))
   

   map_location参数用于指定模型加载到的设备（CPU或GPU）。

3. 验证模型： 加载模型后，你可以使用一些测试数据来验证模型的性能：

   # 假设有一些测试数据
   test_data = torch.randn(1, 10).to(device)
   test_labels = torch.randn(1, 1).to(device)
   
   # 使用模型进行预测
   predictions = model(test_data)
   print(predictions)
   

注意事项

• 设备一致性：在加载模型时，确保模型和数据在同一设备上（CPU或GPU）。
• 版本兼容性：确保保存和加载模型的PyTorch版本一致，否则可能会出现兼容性问题。
• 自定义对象：如果你的模型包含自定义层或函数，你可能需要提供额外的代码来加载这些自定义对象。

通过以上步骤，你可以在Ubuntu上使用PyTorch保存和加载模型。