PyTorch在Linux上如何进行模型保存与加载

在PyTorch中，模型的保存和加载通常使用torch.save()和torch.load()函数。以下是在Linux上进行模型保存与加载的步骤：

模型保存

1. 定义模型： 首先，你需要定义你的模型。

   import torch
   import torch.nn as nn
   
   class MyModel(nn.Module):
       def __init__(self):
           super(MyModel, self).__init__()
           self.fc = nn.Linear(10, 5)
   
       def forward(self, x):
           return self.fc(x)
   
   model = MyModel()
   

2. 训练模型： 训练你的模型，使其达到一定的性能。

3. 保存模型： 使用torch.save()函数保存整个模型或模型的状态字典。

   # 保存整个模型
   torch.save(model, 'model.pth')
   
   # 或者只保存模型的状态字典
   torch.save(model.state_dict(), 'model_state_dict.pth')
   

模型加载

1. 加载模型： 使用torch.load()函数加载模型。如果你之前保存的是整个模型，可以直接加载；如果是状态字典，则需要先创建一个模型实例，然后加载状态字典。

   # 加载整个模型
   loaded_model = torch.load('model.pth')
   
   # 或者加载模型的状态字典
   model = MyModel()  # 创建一个新的模型实例
   model.load_state_dict(torch.load('model_state_dict.pth'))
   

2. 使用模型： 加载完成后，你可以像使用普通模型一样使用它。

   input_data = torch.randn(1, 10)  # 示例输入数据
   output = loaded_model(input_data)
   print(output)
   

注意事项

• 设备兼容性：确保在加载模型时，模型和数据在同一设备上（CPU或GPU）。如果模型是在GPU上训练的，而你在CPU上加载，需要将模型移动到CPU：

  model = model.cpu()
  

• 版本兼容性：不同版本的PyTorch可能会有不同的模型格式，因此在不同版本之间加载模型时可能会遇到问题。确保使用相同或兼容的PyTorch版本。

• 安全性：torch.load()函数在加载模型时可能会执行任意代码，因此只从可信来源加载模型文件。

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