Linux系统下PyTorch的模型保存与加载方法

在Linux系统下，使用PyTorch保存和加载模型的主要方法是使用torch.save()和torch.load()函数。以下是详细的步骤和示例：

保存模型

1. 定义模型： 首先，你需要定义一个PyTorch模型。

   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. 保存模型： 使用torch.save()函数将模型保存到文件中。

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

   这会将整个模型（包括其结构和参数）保存到一个名为model.pth的文件中。

加载模型

1. 加载模型： 使用torch.load()函数从文件中加载模型。

   loaded_model = torch.load('model.pth')
   

2. 使用加载的模型： 加载的模型可以直接用于推理或进一步训练。

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

注意事项

1. 设备兼容性： 如果你在GPU上训练模型，保存的模型会包含GPU相关的信息。在加载模型到CPU时，可以使用map_location参数。

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

2. 版本兼容性： 确保保存和加载模型的PyTorch版本一致，否则可能会出现兼容性问题。

3. 自定义对象： 如果模型中使用了自定义层或函数，需要在加载模型时提供这些自定义对象的定义。

   def custom_function(x):
       return x * 2
   
   torch.save({'model_state_dict': model.state_dict(), 'custom_function': custom_function}, 'model.pth')
   
   # 加载模型时
   checkpoint = torch.load('model.pth')
   model = MyModel()
   model.load_state_dict(checkpoint['model_state_dict'])
   model.custom_function = checkpoint['custom_function']
   

通过以上步骤，你可以在Linux系统下方便地保存和加载PyTorch模型。