1.背景

因為自身的研究方向是多模態目標跟蹤，突然對其他的視覺方向產生了興趣，所以心血來潮的回到最經典的視覺任務手寫數字預測上來，所以這份教程并不是一份非常詳盡的教程，是在一部分pytorch，深度學習基礎上的教程，如果需要的是非常保姆級的教程建議看別的文章

2.準備數據集

這里我才用了直接導torchvision中的dataset包來下載Mnist數據集，也算是一個非常經典的數據集了

# 導入數據集
from torchvision.datasets import MNIST
import torch# 設置隨機種子
torch.manual_seed(3306)# 數據預處理
from torchvision import transforms
# 定義數據轉換
transform = transforms.Compose([transforms.ToTensor(),  # 轉換為 Tensortransforms.Normalize((0.1307,), (0.3081,))  # 標準化
])# 下載 MNIST 數據集
mnist_train = MNIST(root='./dataset_file/mnist_raw', train=True, download=True,transform=transform)
mnist_test = MNIST(root='./dataset_file/mnist_raw', train=False, download=True,transform=transform)
# 查看數據集大小
print(f"MNIST train dataset size: {len(mnist_train)}")
print(f"MNIST test dataset size: {len(mnist_test)}")

其中，MNIST()中的root代表的是數據集存放的位置，download代表是如果當前位置沒有數據集是否需要下載。
transformer則是對數據的處理方式，我這里采用了簡單地轉成tensor和簡單地標準化。

不過這樣子下載下來的數據集是二進制格式的，無法直接查看圖片，當然，如果你需要查看圖片，也有辦法。

# 查看圖片
import matplotlib.pyplot as pltdef show_image(id):img, label = mnist_train[id]img = img.squeeze().numpy()  # 去掉通道維度print(img.shape)# print(img)plt.imshow(img, cmap='gray')plt.title(f"Label: {label}")plt.axis('off')plt.show()show_image(1)

效果

又或者你想要下載的數據集是圖片格式，我這里也準備了代碼

代碼是在別人的基礎上改的，其中數據集存放路徑是dataset_dir，如果需要修改自行打印然后修改位置就好了。

#!/usr/bin/env python3
# -*- encoding utf-8 -*-'''
@File: save_mnist_to_jpg.py
@Date: 2024-08-23
@Author: KRISNAT
@Version: 0.0.0
@Email: ****
@Copyright: (C)Copyright 2024, KRISNAT
@Desc:1. 通過 torchvision.datasets.MNIST 下載、解壓和讀取 MNIST 數據集；2. 使用 PIL.Image.save 將 MNIST 數據集中的灰度圖片以 JPEG 格式保存。
'''import sys, os
sys.path.insert(0, os.getcwd())from torchvision.datasets import MNIST
import PIL
from tqdm import tqdmif __name__ == "__main__":home_dir = os.path.abspath('.')root = os.path.abspath(os.path.join(home_dir, '../dataset_file'))print(root)# exit(0)# 圖片保存路徑dataset_dir = os.path.join(root, 'mnist_jpg')if not os.path.exists(dataset_dir):os.makedirs(dataset_dir)# 從網絡上下載或從本地加載MNIST數據集# 訓練集60K、測試集10K# torchvision.datasets.MNIST接口下載的數據一組元組# 每個元組的結構是: (PIL.Image.Image image model=L size=28x28, 標簽數字 int)training_dataset = MNIST(root='mnist',train=True,download=True,)test_dataset = MNIST(root='mnist',train=False,download=True,)# 保存訓練集圖片with tqdm(total=len(training_dataset), ncols=150) as pro_bar:for idx, (X, y) in enumerate(training_dataset):f = dataset_dir + "/" + "training_" + str(idx) + \"_" + str(training_dataset[idx][1] ) + ".jpg"  # 文件路徑training_dataset[idx][0].save(f)pro_bar.update(n=1)# 保存測試集圖片with tqdm(total=len(test_dataset), ncols=150) as pro_bar:for idx, (X, y) in enumerate(test_dataset):f = dataset_dir + "/" + "test_" + str(idx) + \"_" + str(test_dataset[idx][1] ) + ".jpg"  # 文件路徑test_dataset[idx][0].save(f)pro_bar.update(n=1)

2.定義模型

這里我準備了兩個模型，一個MLP模型和一個簡單地CNN模型，其中MLP模型參數量1M,CNN模型參數量大概8M,當然這倆模型也沒有很仔細的規劃

import torch
import torch.nn as nnclass DigitLinear(nn.Module):def __init__(self):super(DigitLinear, self).__init__()self.fc1 = nn.Linear(28 * 28, 1000)self.fc2 = nn.Linear(1000, 500)self.dropout = nn.Dropout(0.3)self.fc3 = nn.Linear(500, 10)def forward(self, x):x = x.view(-1, 28 * 28)x = self.fc1(x)x = torch.relu(x)x = self.dropout(x)x = self.fc2(x)x = torch.relu(x)x = self.fc3(x)return xclass DigitCNN(nn.Module):def __init__(self):super(DigitCNN,self).__init__()self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)self.fc1 = nn.Linear(64*28*28, 128)self.dropout = nn.Dropout(0.1)self.fc2 = nn.Linear(128, 10)def forward(self, x):# print("x.shape:", x.shape)B,N,H,W = x.shapex = self.conv1(x)x = torch.relu(x)x = self.conv2(x)x = torch.relu(x)x = x.view(B, -1)  # 展平x = self.fc1(x)x = torch.relu(x)x = self.dropout(x)x = self.fc2(x)return x

3.dataloader和訓練

這里的代碼就很簡單了，就是一些參數的選擇，例如epoch，batchsize。其中的訓練函數我寫的買有很全面，只是勉強滿足了訓練功能，還有好多可以優化的點，比如打印fps，斷點續訓練啥的，不過這個任務提不起勁去干這事，大家可以自行優化。

# 數據加載器
from torch.utils.data import DataLoader
from lib.model.DigitModel import DigitLinear,DigitCNN
# 定義數據加載器
batch_size = 256
train_loader = DataLoader(mnist_train, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(mnist_test, batch_size=batch_size, shuffle=False)epoch = 50
# 訓練模型
net = DigitLinear() # 參數量1M 97.50%
# net = DigitCNN() # 參數量8M 98.81%
net.cuda()# 定義損失函數和優化器
import torch.optim as optim
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 訓練函數def train_model(model, train_loader, criterion, optimizer, num_epochs=10):model.train()  # 設置模型為訓練模式for epoch in range(num_epochs):running_loss = 0.0correct = 0total = 0for i, (inputs, labels) in enumerate(train_loader):inputs= inputs.cuda()y = torch.tensor(torch.zeros((inputs.shape[0],10), dtype=torch.float)).cuda()y[torch.arange(inputs.shape[0]), labels] = 1optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, y)loss.backward()optimizer.step()running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels.cuda()).sum().item()epoch_loss = running_loss / len(train_loader)epoch_acc = 100. * correct / totalprint(f'Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')# 訓練模型
train_model(net, train_loader, criterion, optimizer, num_epochs=epoch)

4.訓練模型

有了上面的代碼就可以開始訓練了，我這里訓練的截圖是我的MLP模型，效果不是很好，CNN的效果稍微好一點，比MLP高1%，但是圖忘記截了。反正夠用了，因為本身MNIST的數據就不是很完美，有很多類似于噪聲的數據例如：

這些數字我人眼都分不出是什么玩意。

訓練效果如下

5.測試模型

訓練完當然是測試了
最后我的MLP模型跑了97.50%的準確率

代碼如下

# 測試模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.eval()
correct = 0
total = 0
with torch.no_grad():for inputs, labels in test_loader:inputs, labels = inputs.to(device).float(), labels.to(device).float()outputs = net(inputs)_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels.cuda()).sum().item()# print(f"Predicted: {predicted}, Ground Truth: {targets}")print(f"Accuracy: {correct / total * 100:.4f} %")

6.保存模型

保存模型代碼就更簡單了

# 保存模型
torch.save(net.state_dict(), './digit_model.pth')