Ubuntu下PyTorch訓練腳本怎么寫

在Ubuntu下編寫PyTorch訓練腳本涉及幾個步驟，包括環境設置、數據準備、模型定義、損失函數和優化器的選擇，以及訓練循環的實現。以下是一個簡單的示例，展示了如何編寫一個基本的PyTorch訓練腳本。

步驟1：安裝PyTorch

首先，確保你已經安裝了PyTorch。你可以參考PyTorch官網的安裝指南來安裝適合你系統的版本。

步驟2：編寫訓練腳本

創建一個新的Python文件，例如train.py，并編寫以下代碼：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定義一個簡單的神經網絡
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(-1, 784)  # Flatten the input
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 數據預處理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加載MNIST數據集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# 初始化模型、損失函數和優化器
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 訓練模型
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

print('Finished Training')

代碼解釋

1. 導入必要的庫：

   • torch 和 torch.nn 用于構建和訓練神經網絡。
   • torch.optim 用于定義優化器。
   • torchvision 用于加載和預處理數據。
   • DataLoader 用于批量加載數據。

2. 定義模型：

   • SimpleNet 是一個簡單的全連接神經網絡，包含三個全連接層和一個ReLU激活函數。

3. 數據預處理：

   • 使用 transforms.Compose 組合多個數據預處理操作。
   • ToTensor 將圖像轉換為PyTorch張量。
   • Normalize 對圖像進行標準化。

4. 加載數據集：

   • 使用 datasets.MNIST 加載MNIST數據集，并應用預處理操作。
   • 使用 DataLoader 批量加載數據，并設置 shuffle=True 以打亂數據順序。

5. 初始化模型、損失函數和優化器：

   • SimpleNet 實例化模型。
   • CrossEntropyLoss 作為損失函數。
   • SGD 作為優化器，設置學習率和動量。

6. 訓練模型：

   • 使用嵌套的 for 循環進行訓練。
   • 外層循環遍歷每個epoch。
   • 內層循環遍歷每個batch的數據，計算損失并進行反向傳播和參數更新。
   • 打印每個epoch的損失。

運行腳本

在終端中運行以下命令來執行訓練腳本：

python train.py

這個示例展示了如何在Ubuntu下編寫一個基本的PyTorch訓練腳本。你可以根據具體需求擴展和修改這個腳本，例如添加驗證步驟、保存模型、使用GPU加速等。