使用Pytorch和卷积神经网络进行简单的数字识别（MNIST）

说明：

准备CONDA环境

bash 复制代码

conda create -n py310torch python=3.10
conda activate py310torch
conda install pytorch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install matplotlib tqdm ipywidgets

下载MNIST数据集

Kaggle上有MNIST数据集： https://www.kaggle.com/datasets/hojjatk/mnist-dataset

bash 复制代码

apt-get update
apt-get install curl
mkdir data
cd data
curl -L -o mnist.zip https://www.kaggle.com/api/v1/datasets/download/hojjatk/mnist-dataset
unzip mnist.zip

将MINST数据集保存为图片和标注

借鉴 https://www.kaggle.com/code/hojjatk/read-mnist-dataset 中的读取MNIST数据的代码，并添加保存为图片和标注（按PaddleOCR格式）

python 复制代码

import numpy as np # linear algebra
import struct
from array import array
from os.path  import join

import os
from PIL import Image

#
# MNIST Data Loader Class
#
class MnistDataloader(object):
    def __init__(self, training_images_filepath,training_labels_filepath,
                 test_images_filepath, test_labels_filepath):
        self.training_images_filepath = training_images_filepath
        self.training_labels_filepath = training_labels_filepath
        self.test_images_filepath = test_images_filepath
        self.test_labels_filepath = test_labels_filepath
    
    def read_images_labels(self, images_filepath, labels_filepath):        
        labels = []
        with open(labels_filepath, 'rb') as file:
            magic, size = struct.unpack(">II", file.read(8))
            if magic != 2049:
                raise ValueError('Magic number mismatch, expected 2049, got {}'.format(magic))
            labels = array("B", file.read())        
        
        with open(images_filepath, 'rb') as file:
            magic, size, rows, cols = struct.unpack(">IIII", file.read(16))
            if magic != 2051:
                raise ValueError('Magic number mismatch, expected 2051, got {}'.format(magic))
            image_data = array("B", file.read())        
        images = []
        for i in range(size):
            images.append([0] * rows * cols)
        for i in range(size):
            img = np.array(image_data[i * rows * cols:(i + 1) * rows * cols])
            img = img.reshape(28, 28)
            images[i][:] = img            
        
        return images, labels
            
    def load_data(self):
        x_train, y_train = self.read_images_labels(self.training_images_filepath, self.training_labels_filepath)
        x_test, y_test = self.read_images_labels(self.test_images_filepath, self.test_labels_filepath)
        return (x_train, y_train),(x_test, y_test)

python 复制代码

input_path = './data'
training_images_filepath = join(input_path, 'train-images-idx3-ubyte/train-images-idx3-ubyte')
training_labels_filepath = join(input_path, 'train-labels-idx1-ubyte/train-labels-idx1-ubyte')
test_images_filepath = join(input_path, 't10k-images-idx3-ubyte/t10k-images-idx3-ubyte')
test_labels_filepath = join(input_path, 't10k-labels-idx1-ubyte/t10k-labels-idx1-ubyte')

python 复制代码

mnist_dataloader = MnistDataloader(training_images_filepath, training_labels_filepath, test_images_filepath, test_labels_filepath)
(x_train, y_train), (x_test, y_test) = mnist_dataloader.load_data()

python 复制代码

dest_dir = "./dataset/mnist_train"

os.makedirs(dest_dir, exist_ok = True)

label_lines = []

for i, (image, label) in enumerate(zip(x_train, y_train)):
    fname = f"train_{i}_{label}.png"
    Image.fromarray(np.stack(image)).save(os.path.join(dest_dir, fname))
    label_lines.append(f"mnist_train/{fname}\t{label}")

with open("./dataset/mnist_train_rec.txt", "w") as f:
    f.write("\n".join(label_lines))

python 复制代码

dest_dir = "./dataset/mnist_test"

os.makedirs(dest_dir, exist_ok = True)

label_lines = []

for i, (image, label) in enumerate(zip(x_test, y_test)):
    fname = f"train_{i}_{label}.png"
    Image.fromarray(np.stack(image)).save(os.path.join(dest_dir, fname))
    label_lines.append(f"mnist_test/{fname}\t{label}")

with open("./dataset/mnist_test_rec.txt", "w") as f:
    f.write("\n".join(label_lines))

注意：需要保存为PNG格式，JPG格式保存时会压缩导致和原np.array不一致

准备数据DataLoader

自定义Dataset子类

按照官方教程 https://pytorch.org/tutorials/beginner/basics/data_tutorial.html#creating-a-custom-dataset-for-your-files ，一个自定义的Dataset子类需要实现__init__, __len__ 和 __getitem__这三个方法。

python 复制代码

import os

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
from torch.utils.data import Dataset

import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torchvision.io import decode_image

import numpy as np

from PIL import Image
from tqdm.notebook import trange, tqdm
import time

python 复制代码

class CustomRecDataset(Dataset):
    def __init__(self, data_dir, label_file_or_file_list, transform=None, target_transform=None):
        assert os.path.exists(data_dir), f"data_dir {data_dir} does not exists"
        
        self.data_dir = data_dir
        self.transform = transform
        self.target_transform = target_transform

        self.label_list = []
        
        label_file_list = []
        if isinstance(label_file_or_file_list, str):
            label_file_list.append(label_file_or_file_list)
        for label_file in label_file_list:
            with open(label_file, "r") as f:
                self.label_list.extend(f.readlines())
    
    def __len__(self):
        return len(self.label_list)
    
    def __getitem__(self, idx):
        img_rel_path, label = self.label_list[idx].split("\t")
        label = int(label.replace("\n", ""))
        img_path = os.path.join(self.data_dir, img_rel_path)
        image = Image.open(img_path).numpy()
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

Normalization

通常，我们需要将图片数据进行Normlization处理。

为何要做Normalization：

每个feature具有不同的scale，即有不同的取值范围（比如房子的价格和面积），Normalization 可以移除scale的影响，使得不同的feature对模型的贡献是平等的，防止偏见。
Normalization可以防止因数值过大造成的梯度爆炸或者梯度消失。

关于Normalization，了解更多请参考 https://www.markovml.com/blog/normalization-in-machine-learning

进行Normalization处理需要有MEAN和STD值，CV领域经常使用来自ImageNet的MEAN和STD值。

博客 https://blog.csdn.net/zylooooooooong/article/details/122805833 指出是否使用ImageNet的均值和标准差取决于你的数据：

假设你的数据是"自然场景"的普通照片（人，建筑，动物，不同的照明/角度/背景等等），并且假设你的数据集和 ImageNet 存在类似的偏差（在类别平衡方面），那么使用 ImageNet 的场景统计数据进行规范化就可以了。
如果照片是"特殊的"（颜色过滤，对比度调整，不寻常的光线，等等）或"非自然的主题"（医学图像，卫星地图，手绘等），我建议在模型训练之前正确地规范化你的数据集（计算新的平均值和标准）。

本文不使用来自ImageNet的预训练权重，所以我们需要计算MEAN和STD值：

python 复制代码

ds = CustomRecDataset('dataset', 'dataset/mnist_train_rec.txt')

images = [image for (image, label) in ds]

mean = np.mean(images) / 255.0
std = np.std(images) / 255.0

print(mean, std)

0.1306604762738429 0.30810780385646264

注意：计算MEAN和STD值的时候，只使用训练集的数据，避免泄露来自测试集的数据。

Transforms

python 复制代码

train_transforms = transforms.Compose([
        transforms.RandomRotation(5, fill=(0,)),
        transforms.RandomCrop(28, padding=2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[mean], std=[std])
                    ])

test_transforms = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[mean], std=[std])
                    ])

其中：

RandomRotation随机旋转图片(-x, +x)度
RandomCrop应用padding之后，随机裁剪
ToTensor转换为Tensor并scale到[0.0,1.0]之间

准备训练集，验证集，测试集

python 复制代码

train_data = CustomRecDataset('dataset', 'dataset/mnist_train_rec.txt', transform=train_transforms)

test_data = CustomRecDataset('dataset', 'dataset/mnist_test_rec.txt', transform=test_transforms)

MNIST数据集没有直接提供验证集，我们从训练集里面取出10%来当做验证集。

注意：我们可以从训练集取部分数据当验证集，但是不能从测试集取。

python 复制代码

generator1 = torch.Generator().manual_seed(42)
train_data, valid_data = data.random_split(train_data, [0.9, 0.1], generator=generator1)
valid_data.transform=test_transforms

DataLoader

因为随机梯度下降算法需要，我们将训练集的DataLoader的shuffle为True。

python 复制代码

BATCH_SIZE = 64

train_iterator = data.DataLoader(train_data,
                                 shuffle=True,
                                 batch_size=BATCH_SIZE)

valid_iterator = data.DataLoader(valid_data,
                                 batch_size=BATCH_SIZE)

test_iterator = data.DataLoader(test_data,
                                batch_size=BATCH_SIZE)

定义模型

定义一个模型，参考自 https://www.sciencedirect.com/science/article/pii/S2214579620300502?via%3Dihub#se0070

python 复制代码

class DigitNet(nn.Module):
    def __init__(self, output_dim=10):
        super().__init__()

        self.classifier = nn.Sequential(
            nn.Conv2d(1, 16, 3, 1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, 5, 1),
            nn.MaxPool2d(2, 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 3, 1),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.Linear(128, output_dim)
        )
    
    def forward(self, x):
        logits = self.classifier(x)
        return logits

计算模型参数量：

python 复制代码

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

model = DigitNet(10)
print(f'The model has {count_parameters(model):,} trainable parameters')

0.1306604762738429 0.30810780385646264
The model has 65,674 trainable parameters

训练

python 复制代码

optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = model.to(device)
criterion = criterion.to(device)

def calculate_accuracy(y_pred, y):
    top_pred = y_pred.argmax(1, keepdim=True)
    correct = top_pred.eq(y.view_as(top_pred)).sum()
    acc = correct.float() / y.shape[0]
    return acc

python 复制代码

def train(model, iterator, optimizer, criterion, device):

    epoch_loss = 0
    epoch_acc = 0

    model.train()

    for (x, y) in tqdm(iterator, desc="Training", leave=False):

        x = x.to(device)
        y = y.to(device)

        optimizer.zero_grad()

        y_pred = model(x)

        loss = criterion(y_pred, y)

        acc = calculate_accuracy(y_pred, y)

        loss.backward()

        optimizer.step()

        epoch_loss += loss.item()
        epoch_acc += acc.item()

    return epoch_loss / len(iterator), epoch_acc / len(iterator)

python 复制代码

def evaluate(model, iterator, criterion, device):

    epoch_loss = 0
    epoch_acc = 0

    model.eval()

    with torch.no_grad():

        for (x, y) in tqdm(iterator, desc="Evaluating", leave=False):

            x = x.to(device)
            y = y.to(device)

            y_pred = model(x)

            loss = criterion(y_pred, y)

            acc = calculate_accuracy(y_pred, y)

            epoch_loss += loss.item()
            epoch_acc += acc.item()

    return epoch_loss / len(iterator), epoch_acc / len(iterator)

python 复制代码

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

python 复制代码

EPOCHS = 10

best_valid_loss = float('inf')

for epoch in trange(EPOCHS):

    start_time = time.monotonic()

    train_loss, train_acc = train(model, train_iterator, optimizer, criterion, device)
    valid_loss, valid_acc = evaluate(model, valid_iterator, criterion, device)

    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'tut1-model.pt')

    end_time = time.monotonic()

    epoch_mins, epoch_secs = epoch_time(start_time, end_time)

    print(f'Epoch: {epoch+1:02} | Epoch Time: {epoch_mins}m {epoch_secs}s')
    print(f'\tTrain Loss: {train_loss:.3f} | Train Acc: {train_acc*100:.2f}%')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. Acc: {valid_acc*100:.2f}%')

Epoch: 01 | Epoch Time: 0m 22s
	Train Loss: 0.285 | Train Acc: 91.15%
	 Val. Loss: 0.134 |  Val. Acc: 96.06%
Epoch: 02 | Epoch Time: 0m 21s
	Train Loss: 0.112 | Train Acc: 96.62%
	 Val. Loss: 0.115 |  Val. Acc: 96.63%
Epoch: 03 | Epoch Time: 0m 21s
	Train Loss: 0.089 | Train Acc: 97.26%
	 Val. Loss: 0.085 |  Val. Acc: 97.64%
Epoch: 04 | Epoch Time: 0m 21s
	Train Loss: 0.074 | Train Acc: 97.71%
	 Val. Loss: 0.079 |  Val. Acc: 97.59%
Epoch: 05 | Epoch Time: 0m 19s
	Train Loss: 0.066 | Train Acc: 97.97%
	 Val. Loss: 0.105 |  Val. Acc: 96.66%
Epoch: 06 | Epoch Time: 0m 21s
	Train Loss: 0.061 | Train Acc: 98.06%
	 Val. Loss: 0.073 |  Val. Acc: 97.73%
Epoch: 07 | Epoch Time: 0m 21s
	Train Loss: 0.056 | Train Acc: 98.29%
	 Val. Loss: 0.066 |  Val. Acc: 97.94%
Epoch: 08 | Epoch Time: 0m 21s
	Train Loss: 0.054 | Train Acc: 98.32%
	 Val. Loss: 0.087 |  Val. Acc: 97.23%
Epoch: 09 | Epoch Time: 0m 22s
	Train Loss: 0.051 | Train Acc: 98.40%
	 Val. Loss: 0.065 |  Val. Acc: 98.39%
Epoch: 10 | Epoch Time: 0m 22s
	Train Loss: 0.049 | Train Acc: 98.51%
	 Val. Loss: 0.065 |  Val. Acc: 98.15%