手寫數字識別項目

一、準備數據集

首先我們創建一個卷積模型，訓練的時候就需要一個原始的數據集，那么數據集哪里來？Pytorch官網其實有一些數據集，數據集地址

我們使用到的數據集是MNIST

導入包

import torch
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

使用數據集，所有的官方數據集都繼承 torch.utils.data.Dataset，如果你沒有數據集，那download = True，它會聯網下載到你本地。

# label: 數據集傳入的標簽值
def target_transform(label):return torch.tensor(label)ds = MNIST(root='./data',  # 保存或讀取數據的目錄train=True,  # 是否加載訓練數據集download=False,  # 是否下載數據集transform=ToTensor(),  # 用于轉換圖片的函數# target_transform=target_transform  # 用于轉換標簽的函數target_transform=lambda label: torch.tensor(label)  # 直接匿名函數轉換成張量
)

測試打印數據

print(len(ds))
print(ds[0])
print(ds[0][0].shape)

二、定義模型

簡單的圖像識別模型的套路:卷積 -> 激活 -> 池化 -> … -> 卷積 -> 激活 -> 池化 ->展平 -> 全連接層 -> 激活-> … -> 全連接層輸出,會將圖片縮小的同時增加通道數,當特征圖縮小到 10 以內，就結束卷積過程。之后我們會講到LeNet5模型,這兒我們簡單的定義一個模型進行訓練。

from torch import nn# 卷積激活池化 模塊
class ConvActivatePool(nn.Module):def __init__(self, in_channels, out_channels, kernel_size):super().__init__()# 一般卷積后會選擇讓圖片大小保持不變 進行填充self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding='same')self.relu = nn.ReLU()# 池化在此處提取了特征的同時，讓圖片下采樣了self.pool = nn.MaxPool2d(2)def forward(self, x):x = self.conv(x)x = self.relu(x)y = self.pool(x)return yclass NumberRecognition(nn.Module):def __init__(self):super().__init__()self.cap1 = ConvActivatePool(1, 64, 11)self.cap2 = ConvActivatePool(64, 128, 5)# 分類層self.classifier = nn.Sequential(# 展平nn.Flatten(start_dim=1),# 全連接層nn.Linear(128 * 7 * 7, 2048),nn.ReLU(),nn.Dropout(p=0.3),nn.Linear(2048, 1024),nn.ReLU(),# 輸出結果為 10 分類，所以輸出層全連接輸出 10nn.Linear(1024, 10))# x 形狀 (N, C=1, H=28, W=28)def forward(self, x):x = self.cap1(x)# N x 64 x 14 x 14x = self.cap2(x)# N x 128 x 7 x 7# 圖片縮小到 10 以內，則停止卷積# 調用分類器，對圖片進行分類y = self.classifier(x)return yif __name__ == '__main__':import torchmodel = NumberRecognition()x = torch.rand(16, 1, 28, 28)y = model(x)print(y.shape)

三、模型訓練

3.1 導入依賴庫

import math
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split, Subset
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from model import NumberRecognition

3.2 設備設置（CPU/GPU 自動選擇）

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

3.3 超參數定義

EPOCH = 10          # 訓練輪次：整個訓練集遍歷10次
LR = 1e-2           # 學習率：控制參數更新的步長（1e-2 = 0.01）
BATCH_SIZE = 10     # 批次大小：每次訓練用10個樣本更新一次參數
val_rate = 0.2      # 驗證集比例：從訓練集中劃分20%作為驗證集

3.4數據集準備

1.獲取數據集

ds = MNIST(root='./data',        # 數據集保存路徑（若不存在會自動創建）train=True,           # 加載訓練集（False則加載測試集）download=False,       # 是否自動下載數據集（首次運行需設為True）transform=ToTensor(), # 對圖像的變換：PIL→Tensor（0-1歸一化+維度調整）target_transform=lambda label: torch.tensor(label) # 對標簽的變換：int→Tensor
)

2.劃分訓練集與驗證集

ds_total_len = len(ds)          # 總樣本數：MNIST訓練集共60000個樣本
train_len = int(ds_total_len * (1 - val_rate)) # 訓練集樣本數：60000×0.8=48000
val_len = ds_total_len - train_len             # 驗證集樣本數：60000×0.2=12000
train_ds, val_ds = random_split(ds, [train_len, val_len]) # 隨機劃分

3.創建 DataLoader（按批次加載數據）

# 計算總批次數（向上取整，避免最后一批樣本被丟棄）
train_total_batch = math.ceil(train_len / BATCH_SIZE) # 48000/10=4800批
val_total_batch = math.ceil(val_len / BATCH_SIZE)     # 12000/10=1200批# 訓練集DataLoader
train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True  # 訓練集每次epoch前打亂樣本順序（避免模型記憶樣本順序，提升泛化）
)# 驗證集DataLoader
val_dl = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=True  # 驗證集打亂無意義（僅計算損失），建議設為False以提高效率
)

3.5模型初始化與斷點續訓

# 初始化自定義模型（NumberRecognition在model.py中定義，需確保輸入輸出維度匹配）
model = NumberRecognition()# 嘗試加載歷史模型參數（支持斷點續訓）
try:# 加載參數文件（weights_only=True）state_dict = torch.load('./weights/model.pth', weights_only=True)model.load_state_dict(state_dict) # 將參數加載到模型中print('加載模型參數成功')
except:# 若文件不存在（首次訓練），打印提示print('未找到模型參數')# 將模型遷移到指定設備（CPU/GPU）
model.to(device)

3.6損失函數與優化器定義

# 損失函數：交叉熵損失（適合多分類任務，如MNIST的10類數字）
loss_fn = nn.CrossEntropyLoss()# 優化器：Adam優化器（常用優化器，結合SGD的動量和RMSprop的自適應學習率）
optimizer = torch.optim.Adam(model.parameters(),  # 需優化的參數（模型的所有權重和偏置）lr=LR,               # 學習率（與超參數一致）weight_decay=1e-4    # L2正則化（權重衰減，防止模型參數過大導致過擬合）
)

3.7訓練函數（train ()）

# 全局變量：累計訓練損失和批次數量（用于計算平均損失）
train_total_loss = 0.
train_count = 0def train():global train_total_loss, train_count # 聲明使用全局變量print('開始訓練')model.train() # 將模型設為“訓練模式”（關鍵！啟用Dropout/BatchNorm更新）# 遍歷訓練集DataLoader，每次取一個批次for i, (images, labels) in enumerate(train_dl):# 1. 將數據遷移到指定設備（與模型設備一致）images, labels = images.to(device), labels.to(device)# 2. 清空上一輪的梯度（PyTorch梯度會累加，不清空會導致梯度錯誤）optimizer.zero_grad()# 3. 前向傳播：模型預測輸出y_pred = model(images) # 輸出形狀：(BATCH_SIZE, 10)，每一行是10個類的得分# 4. 計算損失（預測值與真實標簽的差距）loss = loss_fn(y_pred, labels)# 5. 累計損失和批次數量（用于后續計算平均損失）train_total_loss += loss.item() # loss是Tensor，用.item()轉為Python數值train_count += 1# 6. 反向傳播：計算參數梯度（自動微分核心）loss.backward()# 7. 優化器更新參數（根據梯度調整權重和偏置）optimizer.step()# 每100個批次打印一次訓練進度（避免打印過于頻繁）if (i + 1) % 100 == 0:avg_loss = train_total_loss / train_countprint(f'BATCH: [{i + 1}/{train_total_batch}]; loss: {avg_loss:.4f}')# 返回本輪訓練的平均損失（用于epoch結束時打印）return train_total_loss / train_count

3.8驗證函數（valid ()）

def valid():# 局部變量：累計驗證損失和批次數量（每輪驗證重新初始化，避免與訓練混淆）val_total_loss = 0.val_count = 0print('開始驗證')model.eval() # 將模型設為“評估模式”（關鍵！禁用Dropout/BatchNorm更新）# 禁用梯度計算（驗證階段無需反向傳播，節省內存和時間）with torch.no_grad():# 遍歷驗證集DataLoaderfor i, (images, labels) in enumerate(val_dl):# 1. 數據遷移到指定設備images, labels = images.to(device), labels.to(device)# 2. 前向傳播（無梯度計算）y_pred = model(images)# 3. 計算驗證損失loss = loss_fn(y_pred, labels)val_total_loss += loss.item()val_count += 1# 每100個批次打印驗證進度if (i + 1) % 100 == 0:avg_loss = val_total_loss / val_countprint(f'BATCH: [{i + 1}/{val_total_batch}]; loss: {avg_loss:.4f}')# 返回本輪驗證的平均損失return val_total_loss / val_count

3.9主訓練循環（多輪訓練與驗證）

# 遍歷所有訓練輪次
for epoch in range(EPOCH):print(f'\nEPOCH: [{epoch + 1}/{EPOCH}]') # 打印當前輪次（從1開始更直觀）# 1. 訓練本輪并獲取訓練平均損失train_loss = train()# 2. 驗證本輪并獲取驗證平均損失val_loss = valid()# 3. 打印本輪訓練結果print(f'EPOCH END; train loss: {train_loss:.4f}; val loss: {val_loss:.4f}')# 訓練結束后，保存最終模型參數（覆蓋原有文件）
torch.save(model.state_dict(), './weights/model.pth')
print('\n模型參數已保存至 ./weights/model.pth')

四、模型訓練完整代碼

import math
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split, Subset
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from model import NumberRecognitiondevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')EPOCH = 10
LR = 1e-2
BATCH_SIZE = 10
val_rate = 0.2ds = MNIST(root='./data', train=True, download=False, transform=ToTensor(),target_transform=lambda label: torch.tensor(label))ds_total_len = len(ds)
train_len = int(ds_total_len * (1 - val_rate))
val_len = len(ds) - train_len
train_ds, val_ds = random_split(ds, [train_len, val_len])train_total_batch = math.ceil(train_len / BATCH_SIZE)
val_total_batch = math.ceil(val_len / BATCH_SIZE)train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=True)model = NumberRecognition()
try:state_dict = torch.load('./weights/model.pth', weights_only=True)model.load_state_dict(state_dict)print('加載模型參數成功')
except:print('未找到模型參數')model.to(device)loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)train_total_loss = 0.
train_count = 0def train():global train_total_loss, train_countprint('開始訓練')model.train()for i, (images, labels) in enumerate(train_dl):# 3. 將數據放到設備上images, labels = images.to(device), labels.to(device)optimizer.zero_grad()y = model(images)loss = loss_fn(y, labels)train_total_loss += loss.item()train_count += 1loss.backward()optimizer.step()if (i + 1) % 100 == 0:print(f'BATCH: [{i + 1}/{train_total_batch}]; loss: {train_total_loss / train_count}')return train_total_loss / train_countdef valid():val_total_loss = 0.val_count = 0print('開始驗證')model.eval()with torch.no_grad():for i, (images, labels) in enumerate(val_dl):images, labels = images.to(device), labels.to(device)y = model(images)loss = loss_fn(y, labels)val_total_loss += loss.item()val_count += 1if (i + 1) % 100 == 0:print(f'BATCH: [{i + 1}/{val_total_batch}]; loss: {val_total_loss / val_count}')return val_total_loss / val_countfor epoch in range(EPOCH):print(f'EPOCH: [{epoch + 1}/{EPOCH}]')train_loss = train()val_loss = valid()print(f'EPOCH END; train loss: {train_loss}; val loss: {val_loss}')torch.save(model.state_dict(), './weights/model.pth')

五、總結流程

1. 加載 MNIST 公開手寫數字數據集（訓練集）
2. 劃分訓練集與驗證集（用于監控過擬合）
3. 加載自定義的數字識別模型（NumberRecognition），支持斷點續訓（加載歷史參數）
4. 定義訓練 / 驗證流程，使用交叉熵損失和 Adam 優化器訓練模型
5. 訓練完成后保存模型參數，便于后續推理或繼續訓練。