Pytorch(五) 自定义Modul完成MNIST手写数字识别(附查看图片)

引入

机器书别手写数字

• 输入: 数字图片
• 处理: 网络模型
• 输出: 识别结果

最核心的部分当然是网络模型的构建了,会在后文中讲解
首先查看一下MINIST数据集是什么

• 该数据集包含60,000个用于训练的示例和10,000个用于测试的示例。
• 数据集包含了0-9共10类手写数字图片，每张图片都做了尺寸归一化，都是28x28大小的灰度图。
  
  其中的每个数字都是灰度图,类似下图
  
  很明显,其图像为单通道

流程

• 定义超参数 ；
• 构建transforms，主要是对图像做变换；
• 下载、加载数据集 MNIST；
• 补充: 对于MNIST图片的查看
• 构建网络模型；
• 定义训练方法；
• 定义测试方法；
• 开始训练模型，输出预测结果；

实例

1.加载必要的库

# 1.加载必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

2.定义超参数

BATCH_SIZE = 64  # 每批处理的数据
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 用GPU还是CPU
EPOCHS = 10  # 训练数据集的轮次

3.构建pipeline

pipeline = transforms.Compose([
    transforms.ToTensor(),  # 将图片转换为Tensor
    transforms.Normalize((0.1307,), (0.3081,))  # 正则化: 降低模型复杂度
])

注: 0.1307和0.3081为官网提供的数据
transforms主要将图片变换为tensor,并做相应的正则化处理
可以当模型出现过拟合现象的时候 降低模型的复杂度

4. 下载数据集

不用手动去官网下载,可以使用torch中的datasets方便的完成数据下载的需求

train_set = datasets.MNIST('data', train=True, download=True, transform=pipeline)
test_set = datasets.MNIST('data', train=False, download=True, transforms=pipeline)

使用DataLoader加载数据

# 加载数据
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE)

• batch_size 可以将数据切成我们指定的batch_size大小
• shuffle可以起到打乱数据的作用,而测试集没有必要打乱数据,为了减少开销,所以就删去了shffle的步骤

下载完成后,会在当前目录下多出一个data文件夹,

各个gz文件的含义大致如下图所述,分别代表了图片和标签

补充: 对于MNIST图片的查看

很好奇MINIST手写图片是长什么样子的,但是很多教程都直接开始构建网络训练,并没有提到如何查看
这里需要使用 cv2 这个库

pip install opencv-python

下面就以下载下来的train-images-idx3-ubyte文件为例

with open('./data/MNIST/raw/train-images-idx3-ubyte', 'rb') as f:
    file = f.read()
print(file)

直接看的话是一堆byte字节,因此需要用opencv 来转换一下

import cv2
import numpy as np
# 读取字节文件
with open('./data/MNIST/raw/train-images-idx3-ubyte', 'rb') as f:
    file = f.read()
image1 = [int(str(item).encode('ascii'), 16) for item in file[16:16 + 784]]
image1_np = np.array(image1, dtype=np.uint8).reshape(28, 28, 1) # 28 *28 像素的图片 单通道 所以 28 * 28 * 1
print(image1_np.shape) #　(28, 28, 1)
cv2.imwrite('test.jpg', image1_np)　# 保存图片

5.构建网络模型

# 5. 构建网络模型
class Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)  # 1∶灰度图片的通道，10∶输出通道，5∶ kerne
        self.conv2 = nn.Conv2d(10, 20, 3)  # 10∶输入通道，20∶输出通道，3∶Kernel
        self.fc1 = nn.Linear(20 * 10 * 10, 500)  # 20*10*10∶输入通道，500∶输出通道
        self.fc2 = nn.Linear(500, 10)  # 500∶输入通道，10∶输出通道(代表了0-9十个数字的概率)

    def forward(self, x):
        input_size = x.size(0)  # batch_size
        # 卷积+激活+池化
        x = self.conv1(x)  # 输入∶batch*1*28*28，输出∶ batch*10*24*24 ( 28 - 5 + 1 = 24)
        x = F.relu(x)  # 保持shpae不变，输出:batch*10*24*24
        x = F.max_pool2d(x, 2, 2)  # 输入∶ batch*10*24*24 输出∶ batch*10*12*12
        # 卷积+激活
        x = self.conv2(x)  # 输入∶batch*10*12*12输出: batch*20*10*10 ( 12 - 3 +=F.relu ( x)
        x = F.relu(x)

        x = x.view(input_size, -1)  # 拉平
        # 全连接层1
        x = self.fc1(x)
        x = F.relu(x)
        # 全连接层2
        x = self.fc2(x)
        
        output = F.log_softmax(x, dim=1)  # 计算分类后每个数字的概率
        
        return output

6.定义优化器

# 6. 定义优化器
model = Digit().to(DEVICE)
optimizer = optim.Adam(model.parameters())

7.定义训练方法

# 7. 定义训练方法
def train_model(model, device, train_loader, optimizer, epoch):
    # 模型训练
    model.train()
    for batch_index, (data, target) in enumerate(train_loader):
        # 部署到device
        data, target = data.to(device), target.to(device)
        # 梯度初始化为0
        optimizer.zero_grad()
        # 训练后的结果
        output = model(data)
        # 计算损失
        loss = F.cross_entropy(output, target)
        # 反向传播
        loss.backward()
        # 参数优化
        optimizer.step()
        if batch_index % 3000 == 0:
            print(f"Train Epoch : {epoch} Loss : {loss.item():.6f}")

8.定义测试方法

# 8. 定义测试方法
def test_model(model, device, test_loader):
    # 模型验证
    model.eval()
    # 正确率
    correct = 0.0
    # 测试损失
    test_loss = 0.0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target).item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print(f"Test ---- Average loss :{test_loss: .4f}, "
              f"Accuracy :{100.0 * correct / len(test_loader.dataset):.3f}")

交叉熵损失函数 cross_entropy


参考文章 Pytorch常用的交叉熵损失函数CrossEntropyLoss()详解

9.结果

# 9. 方法调用
for epoch in range(1, EPOCHS + 1):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_model(model, DEVICE, test_loader)

官网案例

官网定义的model如下, 采用了2个dropout防止过拟合,
损失函数使用

---

素材来自 轻松学Pytorch

---