计算机视觉读书系列（2）——卷积神经网络

对应《Deep Learning for Vision System》书中的"Convolutional neural networks"一章

大纲脉络与书中一致，后附上个人编写torch框架下代码。

---

1、MLP用于图像分类

当我们预计使用MLP模型来对图像数据进行分析处理时，如分类任务，由于模型结构的原因，我们需要将结构化的图像数据拉伸（flatten）为一维数据，例如将一张28x28的数据（矩阵）拉成784维度的一维数据（向量），随后再将其喂入神经网络，最后得到预测结果。

此类做法有两大致命缺陷：

**++（1）参数量过大++。**由于MLP的结构为全连接，因此其参数量会非常大。

**++（2）损坏结构化数据++。**图像数据中形状、结构类特征对图片理解至关重要，拉扯会损失结构化信息。

相应地，卷积神经网络对应的提出权值共享和感受野来解决对应问题。

2、卷积神经网络（CNN）结构

2.1 总览

回顾机器学习（包括深度学习）的整个架构，其包括：输入数据------预处理------特征提取------解释模型。++显然，MLP在图像任务中的特征提取表现较差（事实上MLP作为最后的分类解释模型表现很好）。因此现在需要替换的是流程中的特征提取模块，也即使用CNN来进行特征提取，使用MLP来进行最终的分类。++

特征提取的最终的结果是特征向量，该特征向量可以作为MLP模型的输入，用于最后的模型分类。

特征提取过程中每经过一层则维度较小，但特征图的数量上升，也即维度减小，深度增加。（把每一个特征图理解为一个特征或一个属性，特征越多表示描述的角度和方面就越多，每一个特征图是聚焦某一个特征，因此其数据维度自然会降低）

特征图：特征图（feature map）是卷积核对图像进行卷积所得到的结果。

2.2 特征提取直观认识

如上图所示，一张一维的图片通过（四次）卷积操作获得了四个特征图片，也就是特征图（feature map）， 所以原始数据从 28x28x1 变成了 20x20x4 (举例）。这20x20x4就是四个特征图（feature map）

2.3 图像分类问题直观认识

以人脸识别为例，卷积神经网络通过识别局部特征来最终实现人脸识别。

3、卷积网络基本成分------卷积、池化、全连接

3.1 卷积层

卷积层通常被用于进行特征学习或特征提取。全连接层通常被用于进行分类。

卷积：卷积层作为一个特征发现窗口、滑过整个图像的像素值，并提取有图像对象种有意义的特征。

卷积核（convolution filter）：卷积核也即CNN网络中的权重， 用来操作图像的算子。其长度通常为方形，随机初始化，并随网络训练而调整（学习）。卷积方式是对应位置相乘相加，最后加上偏差。如下：

在数字图像处理中，卷积核被称为算子，不同类型的算子能对图像进行不同类型的操作，如：边缘锐化算子、边缘提取算子、图像增强算子等。 如下图：

因此，可以理解为，每一个卷积核（算子）能得到一个图像的特征，现在网络需要自主通过学习了解如何提取有用的特征，也即通过学习获得卷积核中的参数，不同的卷积核（算子）所提取的特征不同，这也是为什么将提取的结果称为特征图，因为每一个图代表谋一个特征的提取结果。

感受野：卷积核所操作的区域被称为感受野（receptive field）。

特征图（feature map）：卷积核操作过后的图像被称为特征图。

卷积中四个超参数 ：卷积核个数 （通道数，也即每一层的卷积核的个数）、卷积核尺寸、步长、填充。

（1）卷积核个数。如果在这一层使用N个卷积核，那个下一层的feature map的深度或者通道数也即为N。

（2）卷积核尺寸。通常由2x2、3x3、5x5。小的卷积核用于提取图像中更加精细的特征，较大的卷积核会忽略细节而提取更大维度的特征。（经验主义而言：卷积核的大小最小2x2最大5x5）.

（3）步长。滑动的时候跳跃的像素个数。

（4）填充。对卷积后的图像的填充。

步长和填充的作用：能保证卷积后的图像大小和原始图像的大小一致，方便后续的池化操作计算图像输出大小。

3.2 池化层或降采样

如上所述，每一个卷积核都会生成一幅feature map，那么在卷积模型中数据量就会出现较快的增长，因此需要对数据进行降采样，以及对图像数据进行压缩，以期降低数据量。

池化层主要包括两类，最大池化和平均池化。最大池化操作与卷积操作一致，有一个一定大小的"卷积核"对输入图像进行滑动，同时也伴有步长，只不过卷积核中没有参数，**且输出为卷积核所对应区域中元素最大的值，这也即最大池化。**易知，若池化的卷积核尺寸为2x2，那么最终的数据量将缩减为原来的1/4。

（全局）平均池化是一个更加极端的池化方式，其计算整个通道（channel）或说整张特征图数字的平均值。如下：

++池化层主要用于减少数据的维度降低计算量等++。

（注：现有研究也表明，省去池化层，仅仅使用卷积层一直对图像数据进行数据提取，也能得到较好的效果。另一个方面，随着硬件存储等的发展，降低数据量的急迫性没有这么大，因此现在由很多研究舍弃池化层）

卷积被用于在原始数据中提取图像特征，因此在卷积神经网络中，卷积层结构通常是多个卷积模块（包括卷积和池化）的堆叠，用于连续不断的提取图像中的特征，直到提取到足够小的大小，将其拉伸为一个向量，用作最优预测模型（多为MLP-based）的输出。如下：

3.3 全连接层

全连接层也即多层神经网络，卷积模块所获得的结果传入神经网络中，进行图像分类。

如上图所示，将5x5x40压缩成1000维度的向量，再将这个向量传入MLP中进行分类。

4、使用CNN进行图像分类

4.1 组件一个卷积神经网络

本文给出了一个keras框架下的基于CNN网络的图像分类。

将模型结输出在屏幕如下：

卷积： 模型的输入为灰度图像，其维度为（28，28，1）。第一个卷积模块 conv2d_1，其拥有32个卷积核，其输出维度为（28，28，32）。（下同）

池化：然后经过一个池化层，该池化为尺寸为2的最大池化，无参数，其输出的数据维度（14，14，32），通道数没变，但数据维度标为之前的1/2，总数据变为之前的1/4。（下同）

全连接：首先用flatten将提取好的特征图拉伸，将（7，7，64）拉伸为向量（3136）

4.2 参数量分析

关于卷积中参数量的计算：conv2d_1 层共有32个卷积核，size为3，其处理的图像的通道为1，每个卷积核有一个偏执参数。因此每个卷积核的参数量为：（3x3x1+1）x32=320。

又如，conv2d_2的参数应为：（3x3x32+1）x64=18496。卷积核的通道数与处理对象特征图的通道数一致，如处理（14，14，32）的图像（特征图），则单个卷积核的形状为（3，3，32）。

可训练参数与不可训练参数：在后续的计算机视觉中，或者后续的视觉大模型中，经常会使用预训练大模型。预训练大模型是指在大数据集上训练好的模块，如在ImageNet数据集上训练的特征提取模型，我们可以将其视作图像特征提取的工具，这个工具不会参与我们的模型训练，因为他是训练好的，因此在实际使用时将其视为一个特征提取工具冻结其参数。

5、增加dropout以防止过拟合

5.1 什么是过拟合

欠拟合：模型没能拟合训练数据，通常是模型过于简单。

过拟合：对于训练数据拟合太过，记住了训练数据但没有真的学习到特征，通常是建立了一个比较复杂的网络模型（或者数据量太少）。

5.2 什么是dropout层

dropout层能平衡我们的神经网络，让我们每个神经元节点能被平等的对待，使得预测结果不会过度的依赖于某个特征或某个神经元节点。其工作方式如下：

在实际过程中，每个神经元都有一定的概率（dropout的超参数）被失活（被out），如上图所示，而每个神经元表示一个特征，以达到训练结果不会过度依赖某一个特征的目的。

++dropout减少了神经元之间的相互依赖，在这种情况下，模型一定程度上可能看作多个若分类模型的组合模型。++

5.3 为什么需要dropout层

防止过拟合！！！

6、使用卷积神经网络对彩色图像进行分类

彩色图像为三通道图像，因此在数据输入部分需要采用三通道的卷积核。

使用卷积神经网络对彩色图像进行分类的步骤如下：

（1）加载数据集

（2）图像预处理：

rescale：将图像/255转化为0-1之间数值

准备标签：one-hot编码

数据集分割

定义框架

组建模型：将数据集、模型、优化器等合并

训练模型

加载模型（从最优的记录中）

评估模型

附：个人编写基于torch实现书中示例

所有代码均上传至（数据太大，自行网站下载）：

https://github.com/xurongtang/CV_practicehttps://github.com/xurongtang/CV_practice

附1、手写数字识别项目

脚本1（dataset_load.py） ：加载MINIST数据集，参考博客：Fashion MNIST数据集的处理------"...-idx3-ubyte"文件解析_t10k-images-idx3-ubyte-CSDN博客文章浏览阅读8.4k次，点赞6次，收藏39次。MNIST数据集可能是计算机视觉所接触的第一个图片数据集。而 Fashion MNIST 是在遵循 MNIST 的格式和大小的基础上，提升了一定的难度，在比较算法的性能时可以有更好的区分度。Fashion MNIST 数据集包含 60000 张图片的训练集和 10000 张图片的测试集。图片的大小为 28×28，共784个像素。像素的灰度值介于0~255之间的整数。_t10k-images-idx3-ubytehttps://blog.csdn.net/weixin_43276033/article/details/124163762如下：

import torch
import cv2,os
import numpy as np
import struct
import matplotlib.pyplot as plt

def __decode_idx3_ubyte(file):
    """
    解析数据文件
    """
    # 读取二进制数据
    with open(file, 'rb') as fp:
        bin_data = fp.read()
    
    # 解析文件中的头信息
    # 从文件头部依次读取四个32位，分别为：
    # magic，numImgs, numRows, numCols
    # 偏置
    offset = 0
    # 读取格式: 大端
    fmt_header = '>iiii'
    magic, numImgs, numRows, numCols = struct.unpack_from(fmt_header, bin_data, offset)
    # print(magic,numImgs,numRows,numCols)
    
    # 解析图片数据
    # 偏置掉头文件信息
    offset = struct.calcsize(fmt_header)
    # 读取格式
    fmt_image = '>'+str(numImgs*numRows*numCols)+'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).reshape(numImgs, numRows, numCols)
    return data

def __decode_idx1_ubyte(file):
    """
    解析标签文件
    """
    # 读取二进制数据
    with open(file, 'rb') as fp:
        bin_data = fp.read()
    
    # 解析文件中的头信息
    # 从文件头部依次读取两个个32位，分别为：
    # magic，numImgs
    # 偏置
    offset = 0
    # 读取格式: 大端
    fmt_header = '>ii'
    magic, numImgs = struct.unpack_from(fmt_header, bin_data, offset)
    # print(magic,numImgs)
    # 解析图片数据
    # 偏置掉头文件信息
    offset = struct.calcsize(fmt_header)
    # 读取格式
    fmt_image = '>'+str(numImgs)+'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset))
    return data

def get_MINIST():
    # 文件路径
    data_path = 'E:/ComputerVision_Proj/dataset/MINIST/DATA/'
    file_names = ['t10k-images-idx3-ubyte',
                    't10k-labels-idx1-ubyte',
                    'train-images-idx3-ubyte',
                    'train-labels-idx1-ubyte']
    test_set = (__decode_idx3_ubyte(os.path.join(data_path, file_names[0])),
                __decode_idx1_ubyte(os.path.join(data_path, file_names[1])))
    train_set = (__decode_idx3_ubyte(os.path.join(data_path, file_names[2])),
                __decode_idx1_ubyte(os.path.join(data_path, file_names[3])))
    return train_set, test_set

if __name__ == '__main__':
    train_set, test_set = get_MINIST()
    print(train_set[0].shape,train_set[1].shape)
    print(test_set[0].shape,test_set[1].shape)
    test = np.array(train_set[0][0])
    plt.imshow(test)
    plt.show()

脚本2（cnn_digital_recognize）：完全复现下列结构。

代码如下：

## Author: 撞南墙者 ##
## Date: 2024-12-7 ##
## CopyRight: 随意使用，欢迎交流 ## 
import torch,random
import copy
import numpy as np
import torch.nn as nn
import matplotlib.pyplot as plt
from dataset_load import get_MINIST
from torch.utils.data import DataLoader
from tqdm import tqdm

random.seed(2021)
np.random.seed(2021)
torch.manual_seed(2021)

class Custom_dataset(torch.utils.data.Dataset):
    def __init__(self,dataset):
        self.input_dataset,self.label_dataset = dataset
        self.one_hot_code_flag = True
        self.num_classes = 10

    def __getitem__(self,index):
        input_img = self.input_dataset[index]
        if self.one_hot_code_flag:
            label_img = self.one_hot_code(self.label_dataset[index],self.num_classes)
        else:
            label_img = self.label_dataset[index]
        return input_img,label_img
    
    def __len__(self):
        length = self.input_dataset.shape[0]
        assert length == self.label_dataset.shape[0]
        return length

    @staticmethod
    def one_hot_code(input_data: torch.Tensor,num_classes) -> torch.Tensor:
        one_hot = np.zeros((num_classes),dtype=np.int8)
        for i in range(num_classes):
            if input_data == i:
                one_hot[i] = 1
                break
        one_hot = torch.Tensor(one_hot)
        return one_hot

# 定义模型
class CNN_to_digital_number_recognize(nn.Module):
    
    def __init__(self,input_dim):
        super().__init__()
        self.input_dim = input_dim
        ksize1 = 5
        ksize2 = 5
        padding1 = 0    # ksize1//2
        padding2 = 0    # ksize2//2
        final_channels = 20
        pooling_size = 2
        # 如果想保持图像大小，padding = kernel_size//2
        self.conv1 = nn.Conv2d(in_channels=1,out_channels=10,kernel_size=ksize1,padding=padding1)
        self.pool1 = nn.MaxPool2d(kernel_size=pooling_size)
        self.conv2 = nn.Conv2d(in_channels=10,out_channels=final_channels,kernel_size=ksize2,padding=padding2)
        self.pool2 = nn.MaxPool2d(kernel_size=pooling_size)
        # 计算最后的元素总数
        final_elements = final_channels * ((((input_dim[1]-(ksize1//2)*2) // pooling_size) - (ksize2//2)*2) // pooling_size)**2
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(final_elements,512)
        self.linear2 = nn.Linear(512,10)
        # 构建
        self.ConvBlock1 = nn.Sequential(
            self.conv1,
            nn.ReLU(),
            self.pool1)
        self.ConvBlock2 = nn.Sequential(
            self.conv2,
            nn.ReLU(),
            self.pool2)
        self.LinearBlock = nn.Sequential(
            self.flatten,
            self.linear1,
            nn.ReLU(),
            self.linear2,
            nn.Softmax(dim=1))

    def forward(self,x):
        # 卷积输入的数据维度为 (Batch,channel,height,weight)
        x = x.view(-1,self.input_dim[0],self.input_dim[1],self.input_dim[2]).to(torch.float)
        x = self.ConvBlock1(x)
        x = self.ConvBlock2(x)
        x = self.LinearBlock(x)
        # print(x.shape)
        return x

def training_process(train_set,model,params_dict: dict):
    ## 参数设置 ##
    BATCHSIZE = params_dict['batchsize']
    LEARNING_RATE = params_dict['learning_rate']
    EPOCHS = params_dict['epochs']
    device = torch.device(params_dict['device'])
    val_rate = 0.15
    ## 模型准备 ## 
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(),lr=LEARNING_RATE)
    loss_func = nn.CrossEntropyLoss()  # 对于分类问题
    ## 数据准备 ##
    length = train_set[0].shape[0]
    train_dataset = Custom_dataset((train_set[0][:int(length*(1-val_rate))],train_set[1][:int(length*(1-val_rate))]))
    train_loader = DataLoader(train_dataset,batch_size=BATCHSIZE,shuffle=True)
    val_dataset = Custom_dataset((train_set[0][:int(length*val_rate)],train_set[1][:int(length*val_rate)]))
    val_loader = DataLoader(val_dataset,batch_size=BATCHSIZE,shuffle=True)
    ## 定义模型验证/测试函数 ##
    def val(model,val_loader):
        val_loss = []
        model.eval()
        for _,(batch_x,batch_y) in enumerate(val_loader):
            batch_x = batch_x.float().to(device)
            batch_y = batch_y.float().to(device)
            batch_out = model(batch_x)
            # 计算误差
            loss = loss_func(batch_out,batch_y)
            val_loss.append(loss.item())
        model.train()
        return np.mean(val_loss)
    ## 记录训练过程 ##
    train_loss_ls = []
    val_loss_ls = []
    early_stopping_count = 0
    ## 开始训练 ## 
    print("Start training...")
    for epo in tqdm(range(EPOCHS),colour='yellow'):
        batch_loss = []
        for _,(batch_x,batch_y) in enumerate(train_loader):
            optimizer.zero_grad()
            batch_x = batch_x.float().to(device)
            batch_y = batch_y.float().to(device)
            batch_out = model(batch_x)
            # 计算误差
            loss = loss_func(batch_out,batch_y)
            batch_loss.append(loss.item())
            loss.backward()
            optimizer.step()
            if params_dict['device'] == 'cuda':
                torch.cuda.empty_cache()
        train_loss_ls.append(np.mean(batch_loss))
        val_loss_ls.append(val(model,val_loader))
        if epo == 0:
            best_val_loss = val_loss_ls[-1]
        # print(f"Iters {epo+1:>3}/{EPOCHS} is end, Train_loss:{train_loss_ls[-1]:>8.4f}, Valid_loss:{val_loss_ls[-1]:>8.4f}.")
        if val_loss_ls[-1] <= best_val_loss:
            best_val_loss = val_loss_ls[-1]
            best_model = copy.deepcopy(model)
            early_stopping_count = 0
            # print("Best model has been saved.")
        else:
            early_stopping_count += 1
        if early_stopping_count >= 5:
            print("Early stopping by 5 epochs...")
            break
    print("Training finished.")
    return best_model.eval()

def main():
    train_set,test_set = get_MINIST()
    print("train dataset size: ",train_set[0].shape)
    print("test dataset size: ",test_set[0].shape)
    # print(np.unique(np.array(train_set[1]),return_counts=True))
    Model = CNN_to_digital_number_recognize(input_dim=(1,28,28))
    trained_model = training_process(train_set,
                                     Model,
                                     params_dict={
                                            'batchsize':8,
                                            'learning_rate':0.001,
                                            'epochs':100,
                                            'device':'cuda'
                                            })
    # 简单测试，计算预测准确率
    count = 0
    trained_model.cpu()
    test_input_data,test_label_data = test_set
    print("Testing...")
    for i,input_test in enumerate(tqdm(test_input_data,colour='yellow')):
        res = trained_model(input_test.unsqueeze(0))
        pred_label = torch.argmax(res).item()
        if pred_label == test_label_data[i]:
            count += 1
    print(f"Accuracy:{(count/len(test_input_data))*100:>6.2f}%")
    
if __name__ == '__main__':
    main()

代码运行结果：测试集精度约为98.95%

附2、十类别彩色图像分类（CIFAR-10 and CIFAR-100 datasets）

首先数据读取

在附1的数据处理代码中加入以下函数即可

## Author: 撞南墙者 ##
## Date: 2024-12-7 ##
## CopyRight: 随意使用，欢迎交流 ## 

def get_CIFAR10():
    root_path = 'E:/ComputerVision_Proj/dataset/cifar-10-batches-py/'
    train_path = ['data_batch_1', 'data_batch_2', 'data_batch_3', 'data_batch_4', 'data_batch_5']
    import pickle
    res_data = []
    res_labels = []
    for i in range(5):
        with open(os.path.join(root_path, root_path+train_path[i]), 'rb') as fp:
            res = pickle.load(fp, encoding='bytes')
        labels = res[b'labels']
        data = res[b'data']
        assert len(labels) == len(data)
        dataset_input = np.zeros((len(labels), 3, 32, 32))
        dataset_labels = np.array(labels,dtype=np.uint8)
        for j in range(len(labels)):
            dataset_input[j] = data[j].reshape(3, 32, 32)
        res_data.append(dataset_input)
        res_labels.append(dataset_labels)
    all_train_dataset_input = np.concatenate(res_data, axis=0)
    all_train_dataset_labels = np.concatenate(res_labels, axis=0)
    # test dataset
    with open(os.path.join(root_path, 'test_batch'), 'rb') as fp:
        res = pickle.load(fp, encoding='bytes')
    labels = res[b'labels']
    data = res[b'data']
    assert len(labels) == len(data)
    test_input = np.zeros((len(labels), 3, 32, 32))
    test_labels = np.array(labels,dtype=np.uint8)
    for j in range(len(labels)):
        test_input[j] = data[j].reshape(3, 32, 32)
    # 转化
    train_input = torch.tensor(all_train_dataset_input)
    train_labels = torch.tensor(all_train_dataset_labels)
    test_input = torch.tensor(test_input)
    test_labels = torch.tensor(test_labels)
    return (train_input,train_labels),(test_input,test_labels)

其结果为：

只需在原来代码中改动部分参数即可。

## Author: 撞南墙者 ##
## Date: 2024-12-7 ##
## CopyRight: 随意使用，欢迎交流 ## 
import torch,random
import copy
import numpy as np
import torch.nn as nn
import matplotlib.pyplot as plt
from dataset_load import get_CIFAR10
from torch.utils.data import DataLoader
from tqdm import tqdm

random.seed(2021)
np.random.seed(2021)
torch.manual_seed(2021)

class Custom_dataset(torch.utils.data.Dataset):
    def __init__(self,dataset):
        self.input_dataset,self.label_dataset = dataset
        self.one_hot_code_flag = True
        self.num_classes = 10

    def __getitem__(self,index):
        input_img = self.input_dataset[index]
        if self.one_hot_code_flag:
            label_img = self.one_hot_code(self.label_dataset[index],self.num_classes)
        else:
            label_img = self.label_dataset[index]
        return input_img,label_img
    
    def __len__(self):
        length = self.input_dataset.shape[0]
        assert length == self.label_dataset.shape[0]
        return length

    @staticmethod
    def one_hot_code(input_data: torch.Tensor,num_classes) -> torch.Tensor:
        one_hot = np.zeros((num_classes),dtype=np.int8)
        for i in range(num_classes):
            if input_data == i:
                one_hot[i] = 1
                break
        one_hot = torch.Tensor(one_hot)
        return one_hot

# 定义模型
class CNN_to_digital_number_recognize(nn.Module):
    
    def __init__(self,input_dim):
        super().__init__()
        self.input_dim = input_dim
        ksize1 = 5
        ksize2 = 5
        padding1 = 0    # ksize1//2
        padding2 = 0    # ksize2//2
        final_channels = 20
        pooling_size = 2
        # 如果想保持图像大小，padding = kernel_size//2
        self.conv1 = nn.Conv2d(in_channels=3,out_channels=10,kernel_size=ksize1,padding=padding1)
        self.pool1 = nn.MaxPool2d(kernel_size=pooling_size)
        self.conv2 = nn.Conv2d(in_channels=10,out_channels=final_channels,kernel_size=ksize2,padding=padding2)
        self.pool2 = nn.MaxPool2d(kernel_size=pooling_size)
        # 计算最后的元素总数
        final_elements = final_channels * ((((input_dim[1]-(ksize1//2)*2) // pooling_size) - (ksize2//2)*2) // pooling_size)**2
        self.flatten = nn.Flatten()
        self.linear1 = nn.Linear(final_elements,512)
        self.linear2 = nn.Linear(512,10)
        # 构建
        self.ConvBlock1 = nn.Sequential(
            self.conv1,
            nn.ReLU(),
            nn.Dropout(0.5),
            self.pool1)
        self.ConvBlock2 = nn.Sequential(
            self.conv2,
            nn.ReLU(),
            nn.Dropout(0.5),
            self.pool2)
        self.LinearBlock = nn.Sequential(
            self.flatten,
            self.linear1,
            nn.ReLU(),
            self.linear2,
            nn.Softmax(dim=1))

    def forward(self,x):
        # 卷积输入的数据维度为 (Batch,channel,height,weight)
        x = x.view(-1,self.input_dim[0],self.input_dim[1],self.input_dim[2]).to(torch.float)
        x = self.ConvBlock1(x)
        x = self.ConvBlock2(x)
        x = self.LinearBlock(x)
        # print(x.shape)
        return x

def training_process(train_set,model,params_dict: dict):
    ## 参数设置 ##
    BATCHSIZE = params_dict['batchsize']
    LEARNING_RATE = params_dict['learning_rate']
    EPOCHS = params_dict['epochs']
    device = torch.device(params_dict['device'])
    val_rate = 0.15
    ## 模型准备 ## 
    model.to(device)
    optimizer = torch.optim.SGD(model.parameters(),lr=LEARNING_RATE)
    loss_func = nn.CrossEntropyLoss()  # 对于分类问题
    ## 数据准备 ##
    length = train_set[0].shape[0]
    train_dataset = Custom_dataset((train_set[0][:int(length*(1-val_rate))],train_set[1][:int(length*(1-val_rate))]))
    train_loader = DataLoader(train_dataset,batch_size=BATCHSIZE,shuffle=True)
    val_dataset = Custom_dataset((train_set[0][:int(length*val_rate)],train_set[1][:int(length*val_rate)]))
    val_loader = DataLoader(val_dataset,batch_size=BATCHSIZE,shuffle=True)
    ## 定义模型验证/测试函数 ##
    def val(model,val_loader):
        val_loss = []
        model.eval()
        for _,(batch_x,batch_y) in enumerate(val_loader):
            batch_x = batch_x.float().to(device)
            batch_y = batch_y.float().to(device)
            batch_out = model(batch_x)
            # 计算误差
            loss = loss_func(batch_out,batch_y)
            val_loss.append(loss.item())
        model.train()
        return np.mean(val_loss)
    ## 记录训练过程 ##
    train_loss_ls = []
    val_loss_ls = []
    early_stopping_count = 0
    ## 开始训练 ## 
    print("Start training...")
    for epo in tqdm(range(EPOCHS),colour='yellow'):
        batch_loss = []
        for _,(batch_x,batch_y) in enumerate(train_loader):
            optimizer.zero_grad()
            batch_x = batch_x.float().to(device)
            batch_y = batch_y.float().to(device)
            batch_out = model(batch_x)
            # 计算误差
            loss = loss_func(batch_out,batch_y)
            batch_loss.append(loss.item())
            loss.backward()
            optimizer.step()
            if params_dict['device'] == 'cuda':
                torch.cuda.empty_cache()
        train_loss_ls.append(np.mean(batch_loss))
        val_loss_ls.append(val(model,val_loader))
        if epo == 0:
            best_val_loss = val_loss_ls[-1]
        # print(f"Iters {epo+1:>3}/{EPOCHS} is end, Train_loss:{train_loss_ls[-1]:>8.4f}, Valid_loss:{val_loss_ls[-1]:>8.4f}.")
        if val_loss_ls[-1] <= best_val_loss:
            best_val_loss = val_loss_ls[-1]
            best_model = copy.deepcopy(model)
            early_stopping_count = 0
            # print("Best model has been saved.")
        else:
            early_stopping_count += 1
        if early_stopping_count >= 5:
            print("Early stopping by 5 epochs...")
            break
    print("Training finished.")
    return best_model.eval()

def main():
    train_set,test_set = get_CIFAR10()
    print("train dataset size: ",train_set[0].shape)
    print("test dataset size: ",test_set[0].shape)
    # print(np.unique(np.array(train_set[1]),return_counts=True))
    Model = CNN_to_digital_number_recognize(input_dim=(3,32,32))
    trained_model = training_process(train_set,
                                     Model,
                                     params_dict={
                                            'batchsize':32,
                                            'learning_rate':0.001,
                                            'epochs':100,
                                            'device':'cuda'
                                            })
    # 简单测试，计算预测准确率
    count = 0
    trained_model.cpu()
    test_input_data,test_label_data = test_set
    print("Testing...")
    for i,input_test in enumerate(tqdm(test_input_data,colour='yellow')):
        res = trained_model(input_test.unsqueeze(0))
        pred_label = torch.argmax(res).item()
        if pred_label == test_label_data[i]:
            count += 1
    print(f"Accuracy:{(count/len(test_input_data))*100:>6.2f}%")
    
if __name__ == '__main__':
    main()

结果为，精度在46.10%，可以选择调整参数或者网络架构等提升其精度：

---

欢迎批评指导，共勉！