Day 31 手寫數字辨識實作

引言

前幾天把 PyTorch 的關鍵零件一個個拆開來看過：張量、autograd、Dataset／DataLoader、nn.Module、損失函式與最佳化器。今天要把它們全部組起來，完成深度學習的「經典第一關」：MNIST 手寫數字辨識。MNIST 是由 0 到 9 的灰階手寫數字影像所組成，長期被當作分類模型的入門基準。整個訓練腳本大約只有 80 行，但已涵蓋真實深度學習專案的核心流程，非常適合作為第一個完整作品。

今天會用一個三層的全連接網路（MLP）對付這個任務。預期在 CPU 上跑 5 個 epoch 之後，測試集的正確率可以穩定達到 96% 以上。透過這個範例，可以看到一個真實深度學習專案從資料下載、預處理、訓練到評估的標準流程長什麼樣。

認識 MNIST 資料集

MNIST 資料集由 Yann LeCun 等人在 1998 年整理發布，總共有 70,000 張 28×28 像素的灰階影像：60,000 張訓練影像與 10,000 張測試影像。每張影像對應一個 0 到 9 的標籤，並已經預先切分成訓練與測試兩個集合，因此非常適合作為研究。影像內容是阿拉伯數字的單筆寫法，許多研究人員用它來驗證新的分類模型是否能「正常運作」。

為了避免在每次訓練前都要重新下載，PyTorch 把 MNIST 包進了 torchvision.datasets.MNIST 這個類別。設定 download=True，第一次執行時會自動從官方鏡像站下載約 50 MB 的檔案並解壓縮到 ./data 資料夾；之後再執行就會直接讀取本機檔案。

下載的內容大致包含四個檔案：訓練影像、訓練標籤、測試影像、測試標籤。下載完成後，後續的讀取都在本機完成，不會再連線到外部伺服器。

載入資料與預處理

為了送進全連接網路，每張 28×28 的影像會被攤平成 784 維的向量；像素值原本是 0 到 255，預處理時先用 transforms.ToTensor() 縮放到 0 到 1，再用 Normalize 以均值 0.1307、標準差 0.3081 做正規化。這兩個數字是整個 MNIST 資料集的統計結果，是社群沿用已久的標準值。

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

device = "cuda" if torch.cuda.is_available() else "cpu"

transform = transforms.Compose([
    transforms.ToTensor(),                                    # [0, 255] -> [0, 1]
    transforms.Normalize((0.1307,), (0.3081,)),                # 用 MNIST 的均值／標準差
    transforms.Lambda(lambda x: x.view(-1)),                   # 28x28 -> 784
])

# 第一次執行會自動下載到 ./data，之後直接讀本機
train_dataset = datasets.MNIST(
    root="./data", train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
    root="./data", train=False, download=True, transform=transform
)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

print(f"訓練資料筆數：{len(train_dataset)}")   # 輸出：60000
print(f"測試資料筆數：{len(test_dataset)}")    # 輸出：10000

訓練階段用 batch_size=64 並打亂順序；測試階段通常用更大的 batch（這裡設為 256），因為不需要打亂。Lambda 搭配 view(-1) 把 28×28 的影像攤平成 784 維，是最直覺的做法。如果之後想保留空間結構、改用卷積神經網路（CNN），就不需要這個 Lambda，CNN 會把影像當成 [1, 28, 28] 的張量直接處理，這部分會在 Day 32 與 Day 33 介紹。

建立 MLP 模型

模型本身並不複雜：兩個隱藏層、一個輸出層，激活函式用 ReLU。輸出層有 10 個神經元，對應 0 到 9 共 10 個類別的分數；接下來交給 CrossEntropyLoss，它會在內部自動套用 softmax。

class MLP(nn.Module):
    def __init__(self):
        super().__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 = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)                # 輸出層不接激活函式

model = MLP().to(device)
print(model)

# 輸出：
# MLP(
#   (fc1): Linear(in_features=784, out_features=128, bias=True)
#   (fc2): Linear(in_features=128, out_features=64, bias=True)
#   (fc3): Linear(in_features=64, out_features=10, bias=True)
#   (relu): ReLU()
# )

隱藏層的大小 128 與 64 是經驗值，適合作為示範；如果之後想提高表現，可以嘗試更大的隱藏層或加上 dropout、BatchNorm。這個範例刻意維持簡單，目標是讓讀者看到一個能跑的完整流程。

訓練與評估

把 Day 30 的標準訓練迴圈搬過來，再加上正確率的計算，就是這個專案的核心。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

def evaluate(loader):
    model.eval()
    total, correct, total_loss = 0, 0, 0.0
    with torch.no_grad():
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            logits = model(x)
            total_loss += loss_fn(logits, y).item() * x.size(0)
            correct += (logits.argmax(dim=1) == y).sum().item()
            total += x.size(0)
    return total_loss / total, correct / total

for epoch in range(1, 6):
    model.train()
    running_loss = 0.0
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * x.size(0)

    train_loss = running_loss / len(train_loader.dataset)
    test_loss, test_acc = evaluate(test_loader)
    print(f"epoch {epoch}: train loss = {train_loss:.4f}, "
          f"test loss = {test_loss:.4f}, test acc = {test_acc:.4f}")

把這段程式跑起來，第一次執行會等資料下載，之後每個 epoch 在 CPU 上大約 15 到 30 秒，在 GPU 上會快更多。執行結果大致會像下面這樣（每次數字會略有不同）：

epoch 1: train loss = 0.3014, test loss = 0.1422, test acc = 0.9571
epoch 2: train loss = 0.1187, test loss = 0.0958, test acc = 0.9714
epoch 3: train loss = 0.0771, test loss = 0.0824, test acc = 0.9755
epoch 4: train loss = 0.0536, test loss = 0.0756, test acc = 0.9781
epoch 5: train loss = 0.0398, test loss = 0.0721, test acc = 0.9802

測試正確率從第一個 epoch 的 95.7% 一路爬升到第五個 epoch 的 98% 左右。這個數字不算頂尖（CNN 可以輕鬆達到 99% 以上），但對於純全連接網路已經是合理的成果。從 train loss 持續下降、test loss 同步下降的趨勢看來，模型仍在「學東西」而不是「背答案」。如果想看更細節的評估，例如混淆矩陣，可以參考 Day 39 會介紹的方法。

對單張影像做推論

訓練完成後，最常見的下一步是拿模型去預測新影像。PyTorch 的標準做法是把模型切到 eval() 模式、用 torch.no_grad() 包住推論程式碼，再把單張影像轉成模型期待的形狀。這裡示範從測試集取一張影像、印出預測結果。

model.eval()
x, y = test_dataset[0]                  # 取得第一張測試影像與標籤
x = x.unsqueeze(0).to(device)           # [784] -> [1, 784]，在最前面加上 batch 維度

with torch.no_grad():
    logits = model(x)
    pred = logits.argmax(dim=1).item()

print(f"真實標籤：{y}")                  # 輸出：例如 7
print(f"預測標籤：{pred}")               # 應該與真實標籤相同

unsqueeze(0) 的作用是在最前面加一個 batch 維度，讓單張影像的形狀從 [784] 變成 [1, 784]，符合模型期待的 [batch, features] 結構。如果之後想接受真實世界的影像，只要記得把影像轉成灰階、縮放到 28×28、用同一組 mean/std 做正規化即可。

儲存與載入模型

訓練好的模型應該儲存下來，否則關閉程式後就要重頭再來。PyTorch 提供兩種常見的存法：只存參數（state_dict）或整個模型。只存參數是比較建議的做法，因為這種方式不受類別結構影響，日後就算改了模型定義，只要參數名稱對得起來就能載入。

# 儲存
torch.save(model.state_dict(), "mlp_mnist.pt")

# 載入（先建立模型，再把參數放進去）
model = MLP().to(device)
model.load_state_dict(torch.load("mlp_mnist.pt", map_location=device))
model.eval()

map_location=device 是載入時的關鍵：如果模型是在 GPU 上訓練、但現在要在 CPU 上推論，這個引數會自動把權重搬到 CPU；反之亦然。這在部署或換機器時非常方便。Day 40 會再示範如何把 PyTorch 模型匯出成 ONNX，部署到獨立的推論環境。

結語

今天把整套 PyTorch 訓練流程在 MNIST 上跑了一遍。torchvision.datasets.MNIST 幫我們處理資料下載與讀取，DataLoader 提供批次化的訓練資料，三層 MLP 在 CPU 上就能達到 98% 左右的測試正確率。這是深度學習入門最經典的「Hello World」，之後不管是影像分類、文字分類還是其他任務，幾乎都會沿用同樣的流程：資料 → 模型 → 損失函式 → 最佳化器 → 訓練迴圈 → 評估。如果你已經能獨自完成這個專案，代表已經掌握 PyTorch 最核心的工作流，可以放心進入後面的 CNN、RNN 與遷移學習章節。

明天，我們會進到 CNN 的世界。先理解卷積（convolution）與池化（pooling）為什麼比全連接更適合處理影像，再用手寫範例與小張圖示，看 CNN 怎麼從影像中萃取特徵。

Hao程式筆記

搜尋此網誌