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

在CentOS上使用PyTorch進行模型的保存與加載是一個常見的操作。以下是詳細的步驟和示例代碼，幫助你完成這一任務。

安裝PyTorch

首先，確保你已經安裝了PyTorch。你可以使用以下命令來安裝：

pip install torch torchvision

保存模型

在PyTorch中，你可以使用torch.save()函數來保存模型。以下是一個簡單的示例：

import torch
import torch.nn as nn

# 定義一個簡單的神經網絡
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc(x)
        return x

# 創建模型實例
model = SimpleNet()

# 假設模型已經訓練好
# 保存模型
torch.save(model, 'model.pth')

加載模型

加載模型時，你可以使用torch.load()函數。以下是一個示例：

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

# 如果模型是在不同的設備上訓練的（例如GPU），需要將其移動到當前設備
model.to('cpu')  # 或者 'cuda' 如果你在GPU上工作

檢查模型是否加載成功

你可以通過前向傳播一些數據來檢查模型是否加載成功：

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

# 使用加載的模型進行前向傳播
output = model(input_data)
print(output)

完整示例

以下是一個完整的示例，包括模型的定義、訓練、保存和加載：

import torch
import torch.nn as nn
import torch.optim as optim

# 定義一個簡單的神經網絡
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(784, 10)

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.fc(x)
        return x

# 創建模型實例
model = SimpleNet()

# 定義損失函數和優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 假設我們有一些訓練數據
inputs = torch.randn(64, 1, 28, 28)
labels = torch.randint(0, 10, (64,))

# 訓練模型
for epoch in range(5):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

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

# 加載模型
model = torch.load('model.pth')
model.to('cpu')  # 或者 'cuda' 如果你在GPU上工作

# 檢查模型是否加載成功
output = model(inputs)
print(output)

通過以上步驟，你可以在CentOS上輕松地保存和加載PyTorch模型。希望這些信息對你有所幫助！