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模型。