用 PyTorch 搭建 CNN 實現 MNIST 手寫數字識別

在圖像識別領域，卷積神經網絡（CNN）?憑借其對空間特征的高效提取能力，成為手寫數字識別、人臉識別等任務的首選模型。而 MNIST（手寫數字數據集）作為入門級數據集，幾乎是每個深度學習學習者的 “第一個項目”。

本文將帶大家從零開始，用 PyTorch 搭建一個 CNN 模型完成 MNIST 手寫數字識別任務，不僅會貼出完整代碼，還會逐行解析核心邏輯，幫你搞懂 “每個參數為什么這么設”“每一層的作用是什么”，即使是剛接觸 PyTorch 的新手也能輕松跟上。

一、前置知識與環境準備

在開始前，我們需要先明確兩個核心背景，以及搭建好運行環境：

1. 核心背景速覽

MNIST 數據集：包含 70000 張 28×28 像素的灰度手寫數字圖片（0-9），其中 60000 張為訓練集，10000 張為測試集，每張圖片對應一個 “數字類別” 標簽（0-9）。
CNN 為什么適合？：相比全連接神經網絡，CNN 通過 “卷積層提取局部特征（邊緣、紋理）+ 池化層下采樣”，能大幅減少參數數量、避免過擬合，同時更好地保留圖像的空間結構信息。

2. 環境準備

需要安裝 PyTorch 和 TorchVision（PyTorch 官方的計算機視覺庫，內置 MNIST 數據集）：

# pip安裝命令（根據系統自動匹配版本，若需指定CUDA版本可參考PyTorch官網）
pip install torch torchvision

驗證環境是否安裝成功：

import torch
print(torch.__version__)  # 輸出PyTorch版本，如2.0.1
print(torch.cuda.is_available())  # 輸出True表示支持GPU加速（需NVIDIA顯卡）

二、完整代碼先行（可直接運行）

先貼出完整可運行的代碼，后面會逐段拆解解析：

注意：nn.Sequential()是將網絡層組合在一起，內部不能寫函數

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor# 1. 加載MNIST數據集
train_data = datasets.MNIST(root='data',          # 數據保存路徑train=True,           # 加載訓練集download=True,        # 若路徑下無數據則自動下載transform=ToTensor()  # 將圖像轉為Tensor（0-1歸一化+維度調整：(H,W,C)→(C,H,W)）
)
test_data = datasets.MNIST(root='data',train=False,          # 加載測試集download=True,transform=ToTensor()
)# 2. 數據加載器（分批處理數據）
train_loader = DataLoader(train_data, batch_size=64)  # 每批64個樣本
test_loader = DataLoader(test_data, batch_size=64)# 3. 設備配置（優先GPU，其次CPU）
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
print(f'Using {device} device')  # 打印當前使用的設備# 4. 定義CNN模型
class CNN(nn.Module):def __init__(self):super().__init__()# 卷積塊1：輸入(1,28,28) → 輸出(8,14,14)self.conv1 = nn.Sequential(# 卷積層：1個輸入通道→8個輸出通道，卷積核5×5，步長1，填充2nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5, stride=1, padding=2),nn.ReLU(),  # 激活函數（引入非線性）nn.MaxPool2d(kernel_size=2)  # 池化層：2×2下采樣，尺寸減半)# 卷積塊2：輸入(8,14,14) → 輸出(32,7,7)self.conv2 = nn.Sequential(nn.Conv2d(8, 16, 5, 1, 2),  # 8→16通道，其他參數同上nn.ReLU(),nn.Conv2d(16, 32, 5, 1, 2),  # 16→32通道nn.ReLU(),nn.MaxPool2d(2)  # 下采樣后尺寸14→7)# 卷積塊3：輸入(32,7,7) → 輸出(64,7,7)（無池化，保留尺寸）self.conv3 = nn.Sequential(nn.Conv2d(32, 64, 5, 1, 2),  # 32→64通道nn.ReLU(),nn.Conv2d(64, 64, 5, 1, 2),  # 64→64通道（加深特征提取）nn.ReLU())# 全連接層：輸入(64×7×7) → 輸出10（對應10個數字類別）self.out = nn.Linear(64 * 7 * 7, 10)# 前向傳播（定義數據在模型中的流動路徑）def forward(self, x):x = self.conv1(x)x = self.conv2(x)x = self.conv3(x)x = x.view(x.size(0), -1)  # 展平：(batch_size, 64,7,7) → (batch_size, 64×7×7)output = self.out(x)return output# 5. 初始化模型并移至指定設備
model = CNN().to(device)
print(model)  # 打印模型結構，驗證是否正確# 6. 定義訓練函數
def train(dataloader, model, loss_fn, optimizer):model.train()  # 啟用訓練模式（如BatchNorm、Dropout會生效）batch_count = 1  # 計數批次，用于打印日志for X, y in dataloader:# 將數據移至指定設備（GPU/CPU）X, y = X.to(device), y.to(device)# 前向傳播：計算模型預測值pred = model(X)# 計算損失（多分類任務用CrossEntropyLoss）loss = loss_fn(pred, y)# 反向傳播：更新模型參數optimizer.zero_grad()  # 清空上一輪梯度（避免累積）loss.backward()        # 計算梯度（反向傳播）optimizer.step()       # 根據梯度更新參數（優化器執行）# 每100個批次打印一次損失（監控訓練進度）if batch_count % 100 == 0:loss_value = loss.item()  # 取出損失值（脫離計算圖）print(f'Batch: {batch_count:>4} | Loss: {loss_value:>6.4f}')batch_count += 1# 7. 定義測試函數
def test(dataloader, model, loss_fn):model.eval()  # 啟用評估模式（關閉BatchNorm、Dropout）total_samples = len(dataloader.dataset)  # 測試集總樣本數correct = 0  # 正確預測的樣本數total_loss = 0  # 總損失# 禁用梯度計算（測試階段無需更新參數，節省內存）with torch.no_grad():for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)# 累積損失和正確數total_loss += loss_fn(pred, y).item()# pred.argmax(1)：取每行最大概率的索引（即預測類別），與y比較correct += (pred.argmax(1) == y).type(torch.float).sum().item()# 計算平均損失和準確率avg_loss = total_loss / len(dataloader)  # len(dataloader) = 總批次accuracy = (correct / total_samples) * 100  # 準確率（百分比）print(f'\nTest Result | Accuracy: {accuracy:>5.2f}% | Avg Loss: {avg_loss:>6.4f}\n')# 8. 配置訓練參數并執行
loss_fn = nn.CrossEntropyLoss()  # 多分類交叉熵損失（內置Softmax）
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam優化器，學習率0.001
epochs = 10  # 訓練輪次（整個訓練集遍歷10次）# 循環訓練+測試
for epoch in range(epochs):print(f'=================== Epoch {epoch + 1}/{epochs} ===================')train(train_loader, model, loss_fn, optimizer)  # 訓練一輪test(test_loader, model, loss_fn)  # 測試一輪print("Training Finished!")

三、核心代碼逐段解析

上面的代碼看似長，但邏輯很清晰，我們按 “數據→模型→訓練→測試” 的流程拆解核心部分。

1. 數據加載與預處理

MNIST 數據集的加載全靠torchvision.datasets.MNIST，無需手動下載和解析，非常方便。關鍵參數解析：

root='data'：數據會保存在當前目錄的data文件夾下（自動創建）；
train=True/False：True加載 6 萬張訓練集，False加載 1 萬張測試集；
transform=ToTensor()：這是核心預處理步驟，作用有兩個：
1. 將圖像從 “PIL 格式（0-255 像素值）” 轉為 “Tensor 格式（0-1 歸一化值）”，避免大數值導致梯度爆炸；
2. 調整維度：從圖像默認的(高度H, 寬度W, 通道C)轉為 PyTorch 要求的(通道C, 高度H, 寬度W)（MNIST 是灰度圖，C=1）。

然后用DataLoader將數據集分批：

batch_size=64：每次訓練取 64 個樣本計算梯度（batch_size 越大，訓練越穩定，但內存占用越高）；
DataLoader會自動打亂訓練集（默認shuffle=True），避免模型學習到 “樣本順序” 的無關特征。

2. 設備配置：GPU 加速有多重要？

代碼中這行是 “硬件適配” 的關鍵：

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'

cuda：NVIDIA 顯卡的 GPU 加速（訓練 10 輪可能只需 1-2 分鐘）；
mps：蘋果芯片（M1/M2）的 GPU 加速；
cpu：默認選項（訓練 10 輪可能需要 10-20 分鐘，速度慢很多）。

后續通過model.to(device)和X.to(device)，將模型和數據都移到指定設備上，確保計算在同一設備進行（否則會報錯）。

3. CNN 模型搭建（核心中的核心）

我們定義的CNN類繼承自nn.Module（PyTorch 所有模型的基類），核心是__init__（定義層）和forward（定義數據流動）。

先看模型結構總覽：

輸入(1,28,28) → 卷積塊1 → 輸出(8,14,14) → 卷積塊2 → 輸出(32,7,7) → 卷積塊3 → 輸出(64,7,7) → 展平 → 全連接層 → 輸出(10)

（1）卷積層參數解析

以conv1的第一個卷積層為例：

nn.Conv2d(in_channels=1, out_channels=8, kernel_size=5, stride=1, padding=2)

in_channels=1：輸入通道數（MNIST 是灰度圖，所以 1）；
out_channels=8：輸出通道數 = 卷積核數量（8 個卷積核，提取 8 種不同特征）；
kernel_size=5：卷積核大小（5×5 的窗口，比 3×3 能提取更復雜的局部特征）；
stride=1：卷積核每次滑動 1 個像素（步長越小，特征保留越完整）；
padding=2：填充（在圖像邊緣補 2 個像素），目的是讓卷積后圖像尺寸不變：
👉 尺寸計算公式：輸出尺寸 = (輸入尺寸 - 卷積核尺寸 + 2×padding) / stride + 1
👉 代入：(28 - 5 + 2×2)/1 + 1 = 28，所以卷積后還是 28×28。

（2）激活函數 ReLU

每個卷積層后都加nn.ReLU()，作用是引入非線性：

沒有激活函數的話，多個卷積層疊加還是線性變換，無法擬合復雜數據；
ReLU 的公式：ReLU(x) = max(0, x)，計算簡單、梯度不易消失，是目前最常用的激活函數。

（3）池化層 MaxPool2d

nn.MaxPool2d(kernel_size=2)是 2×2 最大池化，作用是下采樣：

尺寸減半：28×28→14×14，14×14→7×7，大幅減少后續計算量；
保留關鍵特征：取 2×2 窗口的最大值，相當于 “強化局部最顯著的特征”，提高模型魯棒性。

（4）展平與全連接層

卷積塊 3 輸出的是(batch_size, 64, 7, 7)的張量（batch_size 是每批樣本數），需要用x.view(x.size(0), -1)展平為(batch_size, 64×7×7)的一維向量，才能輸入全連接層：

x.size(0)：獲取 batch_size（確保展平后每一行對應一個樣本）；
-1：讓 PyTorch 自動計算剩余維度（64×7×7=3136）；
全連接層nn.Linear(3136, 10)：將 3136 維特征映射到 10 維（對應 0-9 的 10 個類別）。

4. 訓練函數：模型如何 “學習”？

訓練的核心是 “前向傳播算損失→反向傳播求梯度→優化器更新參數” 的循環：

model.train()：啟用訓練模式（比如如果模型有 BatchNorm，會計算當前批次的均值和方差）；
前向傳播：pred = model(X)，用當前模型參數計算預測值；
計算損失：loss = loss_fn(pred, y)，用CrossEntropyLoss（多分類任務專用，內置了 Softmax，無需手動在模型輸出加 Softmax）；
反向傳播：
- optimizer.zero_grad()：清空上一輪的梯度（如果不清空，梯度會累積，導致參數更新錯誤）；
- loss.backward()：自動計算所有可訓練參數的梯度（PyTorch 的自動微分機制）；
- optimizer.step()：用計算出的梯度更新參數（Adam 優化器會自適應調整學習率，比 SGD 收斂更快）。

5. 測試函數：模型學得怎么樣？

測試階段不需要更新參數，核心是計算 “準確率” 和 “平均損失”：

model.eval()：啟用評估模式（關閉 BatchNorm 的批次統計更新、關閉 Dropout）；
with torch.no_grad()：禁用梯度計算（節省內存，加速測試）；
準確率計算：pred.argmax(1) == y，比較預測類別和真實類別，求和后除以總樣本數。

四、預期結果與優化方向

1. 預期訓練結果

在 GPU 上訓練 10 輪后，通常能達到：

測試準確率：98.5% 以上（甚至 99%）；
測試平均損失：0.04 以下。

訓練過程中，損失會逐漸下降，準確率會逐漸上升（如果出現損失不下降或準確率波動，可能是學習率太大或 batch_size 太小）。

2. 模型優化方向

如果想進一步提升性能，可以嘗試這些改進：

增加 Dropout 層：在卷積層或全連接層后加nn.Dropout(0.2)，隨機 “關閉” 20% 的神經元，防止過擬合；
使用學習率調度：比如torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)，每 5 輪將學習率減半，后期精細調整；
加深網絡：增加卷積塊數量（比如再加一個 conv4），或增加每個卷積層的輸出通道數；

數據增強：用torchvision.transforms添加旋轉、平移、縮放等操作，比如：

transform = transforms.Compose([transforms.RandomRotation(5),  # 隨機旋轉±5度transforms.ToTensor()
])

????????數據增強能讓模型看到更多 “變種” 樣本，提升泛化能力。

五、總結

本文用 PyTorch 實現了一個基礎的 CNN 模型，完成了 MNIST 手寫數字識別任務，核心收獲包括：

掌握了 PyTorch 加載數據集、搭建 CNN 模型的基本流程；
理解了卷積層、池化層、激活函數的作用和參數意義；
熟悉了 “訓練 - 測試” 的循環邏輯，以及 GPU 加速的配置方法。

MNIST 是入門任務，但 CNN 的核心思想（特征提取 + 下采樣）可以遷移到更復雜的圖像任務（如 CIFAR-10、ImageNet）。建議大家動手修改代碼，比如調整卷積核大小、學習率、網絡層數，觀察結果變化，這樣才能真正理解每個參數的影響～