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