PyTorch Linux环境下如何进行模型保存与加载

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

模型保存

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

   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)
   

2. 实例化模型： 创建模型的一个实例。

   model = MyModel()
   

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_state_dict = torch.load('model_state_dict.pth')
   loaded_model = MyModel()  # 创建一个新的模型实例
   loaded_model.load_state_dict(model_state_dict)  # 加载状态字典到模型实例
   

注意事项

• 设备兼容性：如果你在一个设备上训练模型（例如GPU），而在另一个设备上加载模型（例如CPU），可能会遇到设备不匹配的问题。可以使用map_location参数来指定加载模型时的设备：

  # 从GPU加载到CPU
  model = torch.load('model.pth', map_location=torch.device('cpu'))
  
  # 或者从CPU加载到GPU（假设你有可用的GPU）
  model = torch.load('model.pth', map_location=torch.device('cuda'))
  

• 版本兼容性：确保保存和加载模型的PyTorch版本相同或兼容，否则可能会遇到不兼容的问题。

• 安全性：torch.load()函数在加载不受信任的文件时可能存在安全风险，因为它会执行文件中的代码。确保只加载来自可信来源的模型文件。

通过以上步骤，你可以在Linux环境下轻松地进行PyTorch模型的保存与加载。