Ubuntu PyTorch如何進行模型保存

在Ubuntu上使用PyTorch進行模型保存非常簡單。以下是保存和加載模型的基本步驟：

保存模型

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')

這會將模型的所有參數和結構保存到一個文件中。

加載模型

1. 加載模型：使用torch.load()函數來加載模型。

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

2. 使用模型：加載后的模型可以直接用于推理或繼續訓練。

input_data = torch.randn(1, 10)  # 示例輸入數據
output = loaded_model(input_data)
print(output)

注意事項

• 設備兼容性：如果你在不同的設備（CPU/GPU）上保存和加載模型，可能會遇到設備不匹配的問題?？梢允褂?code>map_location參數來指定加載模型時的設備。

# 保存模型時指定設備為CPU
torch.save(model.cpu(), 'model.pth')

# 加載模型時指定設備為CPU
loaded_model = torch.load('model.pth', map_location=torch.device('cpu'))

• 版本兼容性：不同版本的PyTorch可能會有不同的模型保存格式。確保在保存和加載模型時使用相同版本的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()

# 保存模型
torch.save(model, 'model.pth')

# 加載模型
loaded_model = torch.load('model.pth', map_location=torch.device('cpu'))

# 使用加載的模型進行推理
input_data = torch.randn(1, 10)
output = loaded_model(input_data)
print(output)

通過以上步驟，你可以在Ubuntu上輕松地保存和加載PyTorch模型。