PyTorch在Ubuntu上的模型保存與加載方法

在Ubuntu上使用PyTorch進行模型的保存與加載，可以按照以下步驟操作：

模型保存

1. 定義模型： 首先，你需要定義一個PyTorch模型。例如：

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. 實例化模型并訓練： 接下來，實例化模型并進行訓練。這里我們假設你已經完成了訓練過程。
3. 保存模型： 使用torch.save()函數來保存整個模型或模型的狀態字典。

model = MyModel()
# 假設模型已經訓練完成
torch.save(model, 'model.pth')  # 保存整個模型
# 或者
torch.save(model.state_dict(), 'model_state_dict.pth')  # 只保存模型的狀態字典

模型加載

1. 加載模型： 使用torch.load()函數來加載模型。如果你之前保存了整個模型，可以直接加載；如果只保存了狀態字典，則需要先實例化模型，然后再加載狀態字典。

# 加載整個模型
model = torch.load('model.pth')

# 或者，如果只保存了狀態字典
model = MyModel()  # 先實例化模型
model.load_state_dict(torch.load('model_state_dict.pth'))

注意：在加載模型時，確保模型類（在本例中為MyModel）已經在當前環境中定義。

2. 使用模型： 加載模型后，你可以像平常一樣使用它進行推理或繼續訓練。

# 假設我們有一些輸入數據
input_data = torch.randn(1, 10)

# 使用模型進行推理
output = model(input_data)
print(output)

注意事項

• 在保存和加載模型時，確保PyTorch的版本一致，以避免潛在的兼容性問題。
• 如果你在不同的環境中保存和加載模型（例如，在Ubuntu上訓練，在Windows上加載），可能需要考慮模型文件路徑的差異以及可能的操作系統兼容性問題。
• 為了確保模型的可移植性，建議將模型保存為狀態字典的形式，并在加載時重新實例化模型類。