文章目錄
- 引言
- 1. 環境準備和數據加載
- 1.1 下載MNIST數據集
- 1.2 數據可視化
- 2. 數據預處理
- 3. 設備配置
- 4. 構建卷積神經網絡模型
- 5. 訓練和測試函數
- 5.1 訓練函數
- 5.2 測試函數
- 6. 模型訓練和評估
- 6.1 初始化損失函數和優化器
- 6.2 訓練過程
- 7. 關鍵點解析
- 8. 完整代碼
- 9. 總結
引言
手寫數字識別是計算機視覺和深度學習領域的經典入門項目。本文將詳細介紹如何使用PyTorch框架構建一個卷積神經網絡(CNN)來實現MNIST手寫數字識別任務。我們將從數據加載、模型構建到訓練和測試，一步步解析整個過程。
1. 環境準備和數據加載
首先，我們需要導入必要的PyTorch模塊：
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
1.1 下載MNIST數據集
MNIST數據集包含60,000個訓練樣本和10,000個測試樣本，每個樣本都是一個28x28像素的灰度手寫數字圖像。
# 下載訓練數據集
training_data = datasets.MNIST(root="data",train=True,download=True,transform=ToTensor(),
)# 下載測試數據集
test_data = datasets.MNIST(root="data",train=False,download=True,transform=ToTensor(),
)
1.2 數據可視化
我們可以使用matplotlib庫來查看數據集中的一些樣本：
from matplotlib import pyplot as pltfigure = plt.figure()
for i in range(9):img, label = training_data[i+59000] # 提取后幾張圖片figure.add_subplot(3,3,i+1)plt.title(label)plt.axis("off")plt.imshow(img.squeeze(), cmap="gray")
plt.show()
2. 數據預處理
為了高效訓練模型，我們需要使用DataLoader將數據集分批次加載：
train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
3. 設備配置
PyTorch支持在CPU、NVIDIA GPU和蘋果M系列芯片上運行，我們可以自動檢測最佳可用設備：
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(CNN, self).__init__()self.conv1 = nn.Sequential(nn.Conv2d(1, 8, 3, 1, 1), # (8,28,28)nn.ReLU(),nn.MaxPool2d(2), # (8,14,14))self.conv2 = nn.Sequential(nn.Conv2d(8, 16, 3, 1, 1), # (16,14,14)nn.ReLU(),nn.MaxPool2d(2), # (16,7,7))self.out = nn.Linear(16*7*7, 10)def forward(self, x):x = self.conv1(x)x = self.conv2(x)x = x.view(x.size(0), -1) # flatten操作output = self.out(x)return outputmodel = CNN().to(device)
這個CNN模型包含：
- 兩個卷積層，每個卷積層后接ReLU激活函數和最大池化層
- 一個全連接輸出層
- 輸入大小：(1,28,28)
- 輸出大小：10（對應0-9的數字類別）
5. 訓練和測試函數
5.1 訓練函數
def train(dataloader, model, loss_fn, optimizer):model.train()batch_size_num = 1for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)loss = loss_fn(pred, y)optimizer.zero_grad()loss.backward()optimizer.step()if batch_size_num % 100 == 0:print(f"loss: {loss.item():>7f} [number:{batch_size_num}]")batch_size_num += 1
5.2 測試函數
def Test(dataloader, model, loss_fn):size = len(dataloader.dataset)num_batches = len(dataloader)model.eval()test_loss, correct = 0, 0with torch.no_grad():for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model(X)test_loss += loss_fn(pred, y).item()correct += (pred.argmax(1) == y).type(torch.float).sum().item()test_loss /= num_batchescorrect /= sizeprint(f"Test result: \n Accuracy:{(100*correct)}%, Avg loss:{test_loss}")
6. 模型訓練和評估
6.1 初始化損失函數和優化器
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
6.2 訓練過程
# 初始訓練和測試
train(train_dataloader, model, loss_fn, optimizer)
Test(test_dataloader, model, loss_fn)# 多輪訓練
epochs = 10
for t in range(epochs):print(f"epoch {t+1}\n---------------")train(train_dataloader, model, loss_fn, optimizer)
print("Done!")# 最終測試
Test(test_dataloader, model, loss_fn)
7. 關鍵點解析
-
數據轉換：使用
ToTensor()將圖像數據轉換為PyTorch張量，并自動歸一化到[0,1]范圍。 -
批處理：DataLoader的
batch_size參數控制每次訓練使用的樣本數量，影響內存使用和訓練速度。 -
模型結構：
- 卷積層提取空間特征
- ReLU激活函數引入非線性
- 最大池化層降低特征圖尺寸
- 全連接層輸出分類結果
-
訓練模式切換：
model.train()和model.eval()分別用于訓練和測試階段，影響某些層(如Dropout和BatchNorm)的行為。 -
優化過程：Adam優化器結合了動量法和自適應學習率的優點，通常能獲得較好的訓練效果。
8. 完整代碼
import torch
from torch import nn #導入神經網絡模塊
from torch.utils.data import DataLoader #數據包管理工具，打包數據
from torchvision import datasets #封裝了很多與圖像相關的模型，數據集
from torchvision.transforms import ToTensor #數據轉換，張量，將其他類型的數據轉換為tensor張量，numpy array'''下載訓練數據集（包含訓練圖片+標簽）'''
training_data = datasets.MNIST( #跳轉到函數的內部源代碼，pycharm按下ctrl + 鼠標點擊root="data", #表示下載的手寫數字 到哪個路徑。60000train=True, #讀取下載后的數據中的訓練集download=True, #如果你之前已經下載過了，就不用下載transform=ToTensor(), #張量，圖片是不能直接傳入神經網絡模型) #對于pytorch庫能夠識別的數據一般是tensor張量'''下載測試數據集（包含訓練圖片+標簽）'''
test_data = datasets.MNIST( #跳轉到函數的內部源代碼，pycharm按下ctrl + 鼠標點擊root="data", #表示下載的手寫數字 到哪個路徑。60000train=False, #讀取下載后的數據中的訓練集download=True, #如果你之前已經下載過了，就不用下載transform=ToTensor(), #Tensor是在深度學習中提出并廣泛應用的數據類型) #Numpy數組只能在CPU上運行。Tensor可以在GPU上運行。這在深度學習應用中可以顯著提高計算速度。
print(len(training_data))'''展示手寫數字圖片，把訓練集中的59000張圖片展示'''
from matplotlib import pyplot as plt
figure = plt.figure()
for i in range(9):img,label = training_data[i+59000] #提取第59000張圖片figure.add_subplot(3,3,i+1) #圖像窗口中創建多個小窗口，小窗口用于顯示圖片plt.title(label)plt.axis("off") #plt.show(I) 顯示矢量plt.imshow(img.squeeze(),cmap="gray") #plt.imshow()將Numpy數組data中的數據顯示為圖像，并在圖形窗口中顯示a = img.squeeze() #img.squeeze()從張量img中去掉維度為1的，如果該維度的大小不為1，則張量不會改變
plt.show()'''創建數據DataLoader（數據加載器）'''
# batch_size:將數據集分為多份，每一份為batch_size個數據
# 優點：可以減少內存的使用，提高訓練速度train_dataloader = DataLoader(training_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)'''判斷當前設備是否支持GPU，其中mps是蘋果m系列芯片的GPU'''
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device") #字符串的格式化，CUDA驅動軟件的功能：pytorch能夠去執行cuda的命令
# 神經網絡的模型也需要傳入到GPU，1個batch_size的數據集也需要傳入到GPU，才可以進行訓練''' 定義神經網絡 類的繼承這種方式'''
class CNN(nn.Module): #通過調用類的形式來使用神經網絡，神經網絡的模型，nn.mdouledef __init__(self): #輸入大小：(1,28,28)super(CNN,self).__init__() #初始化父類self.conv1 = nn.Sequential( #將多個層組合成一起，創建了一個容器，將多個網絡組合在一起nn.Conv2d( # 2d一般用于圖像，3d用于視頻數據（多一個時間維度），1d一般用于結構化的序列數據in_channels=1, # 圖像通道個數，1表示灰度圖（確定了卷積核 組中的個數）out_channels=8, # 要得到多少個特征圖，卷積核的個數kernel_size=3, # 卷積核大小 3×3stride=1, # 步長padding=1, # 一般希望卷積核處理后的結果大小與處理前的數據大小相同，效果會比較好), # 輸出的特征圖為(8,28,28)nn.ReLU(), # Relu層，不會改變特征圖的大小nn.MaxPool2d(kernel_size=2), # 進行池化操作(2×2操作),輸出結果為(8,14,14))self.conv2 = nn.Sequential(nn.Conv2d(8,16,3,1,1), #輸出(16,14,14)nn.ReLU(), #Relu層 (16,14,14)nn.MaxPool2d(kernel_size=2), #池化層，輸出結果為(16,7,7))self.out = nn.Linear(16*7*7,10) # 全連接層得到的結果def forward(self,x): #前向傳播，你得告訴它 數據的流向 是神經網絡層連接起來，函數名稱不能改x = self.conv1(x)x = self.conv2(x)x = x.view(x.size(0),-1) # flatten操作，結果為：(batch_size,64 * 7 * 7)output = self.out(x)return output
model = CNN().to(device) #把剛剛創建的模型傳入到GPU
print(model)def train(dataloader,model,loss_fn,optimizer):model.train() #告訴模型，我要開始訓練，模型中w進行隨機化操作，已經更新w，在訓練過程中，w會被修改的
# pytorch提供2種方式來切換訓練和測試的模式，分別是：model.train() 和 mdoel.eval()
# 一般用法是：在訓練開始之前寫上model.train(),在測試時寫上model.eval()batch_size_num = 1for X,y in dataloader: #其中batch為每一個數據的編號X,y = X.to(device),y.to(device) #把訓練數據集和標簽傳入cpu或GPUpred = model.forward(X) # .forward可以被省略，父類種已經對此功能進行了設置loss = loss_fn(pred,y) # 通過交叉熵損失函數計算損失值loss# Backpropagation 進來一個batch的數據，計算一次梯度，更新一次網絡optimizer.zero_grad() # 梯度值清零loss.backward() # 反向傳播計算得到每個參數的梯度值woptimizer.step() # 根據梯度更新網絡w參數loss_value = loss.item() # 從tensor數據種提取數據出來，tensor獲取損失值if batch_size_num %100 ==0:print(f"loss: {loss_value:>7f} [number:{batch_size_num}]")batch_size_num += 1def Test(dataloader,model,loss_fn):size = len(dataloader.dataset) #10000num_batches = len(dataloader) # 打包的數量model.eval() #測試，w就不能再更新test_loss,correct =0,0with torch.no_grad(): #一個上下文管理器，關閉梯度計算。當你確認不會調用Tensor.backward()的時候for X,y in dataloader:X,y = X.to(device),y.to(device)pred = model.forward(X)test_loss += loss_fn(pred,y).item() #test_loss是會自動累加每一個批次的損失值correct += (pred.argmax(1) == y).type(torch.float).sum().item()a = (pred.argmax(1) == y) #dim=1表示每一行中的最大值對應的索引號，dim=0表示每一列中的最大值對應的索引號b = (pred.argmax(1) == y).type(torch.float)test_loss /= num_batches #能來衡量模型測試的好壞correct /= size #平均的正確率print(f"Test result: \n Accuracy:{(100*correct)}%, Avg loss:{test_loss}")loss_fn = nn.CrossEntropyLoss() #創建交叉熵損失函數對象，因為手寫字識別一共有十種數字，輸出會有10個結果
#
optimizer = torch.optim.Adam(model.parameters(),lr=0.01) #創建一個優化器，SGD為隨機梯度下降算法
# # params：要訓練的參數，一般我們傳入的都是model.parameters()
# # lr:learning_rate學習率，也就是步長
#
# # loss表示模型訓練后的輸出結果與樣本標簽的差距。如果差距越小，就表示模型訓練越好，越逼近真實的模型
train(train_dataloader,model,loss_fn,optimizer) #訓練1次完整的數據。多輪訓練
Test(test_dataloader,model,loss_fn)epochs = 10
for t in range(epochs):print(f"epoch {t+1}\n---------------")train(train_dataloader,model,loss_fn,optimizer)
print("Done!")
Test(test_dataloader,model,loss_fn)
9. 總結
通過本文，我們學習了如何使用PyTorch實現一個完整的手寫數字識別項目。從數據加載、模型構建到訓練和評估，每個步驟都展示了PyTorch框架的簡潔和強大。這個簡單的CNN模型在MNIST數據集上可以達到很高的準確率，為進一步學習更復雜的計算機視覺任務打下了良好基礎。
未來可以嘗試：
- 調整網絡結構(增加層數、改變通道數)
- 嘗試不同的優化器和學習率
- 添加數據增強技術
- 在更復雜的數據集上應用類似方法
希望這篇教程能幫助你入門PyTorch和計算機視覺領域！
使用SASL Plaintext進行身份驗證)
)
![[調試技巧]VS Code如何在代理模式下使用 MCP 工具？](http://pic.xiahunao.cn/[調試技巧]VS Code如何在代理模式下使用 MCP 工具？)