pytorch搭建CNN，手写数字识别，初识神经网络

大概内容：

• torchvision.datasets加载MNIST数据集，及显示里面的内容
• 搭建3层神经网络，介绍相关参数
• 训练网络和测试准确率，图像显示loss的变化
• 使用GPU，即cuda()
• 神经网络模型参数的保存和提取

1.搭建神经网络

一些概念：

• 优化器：常见SDG，RMSporp，Adam等（Adam需要大显存）
• 卷积神经网络：常用在图片识别，视频分析、自然语言处理
• 结构：神经元构成神经层，神经层构成神经网络
• 每个神经层都有其输入和输出，当输入是图片时（以彩色图片为例），输入的是一个三维向量：对应长宽高（像素点和通道数（RGB3通道））
• 卷积：不在对每个单独特征点（像素）输入信息进行处理，而是对一小块的像素区域进行处理，加强了信息的连续性
• 卷积神经网络：有一个批量过滤器，不断在图片上移动，收集信息，不断重复这个过程。一次收集之后，长和宽更小、高更大的图片
• 为什么高变高了呢？这是卷积核定义的，可以定义高度，这个定义值表示这个卷积核在这个区域提取了几次信息
• 这个批量过滤器，也叫滤波器，也叫过滤器，也叫卷积核
• 池化：处理卷积的信息，改变图片的大小

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        # 第一层：输入层
        self.conv1 = nn.Sequential(  # conv1是一个卷积层（一般包括：卷积层，激励函数，池化层），Sequential是其内部的层包裹起来（是一个序贯模型）
            nn.Conv2d(  # 卷积层（2d:二维卷积）：是一个3维过滤器filter，有长宽高，高度指filter的个数（用于提取特征属性）
                in_channels=1, # 输入数据的高度（是上一层的高度，也称卷积核的通道数），图片RGB为3层
                out_channels=16, # 输出的个数（也是filter的个数，是卷积核的个数），是自定义的，表示同时有16个filter在收集信息，这里是使图片变厚的原因
                kernel_size=5, # 表示卷积核的长和宽都是5个像素点
                stride=1,  # 步长：每次卷积核移动的像素点个数(表示每隔多少步跳一下，扫描下一块区域)
                padding=2,  # padding=(kernel_size-stride)/2.表示在图像的边缘填充2个像素点（补全边缘像素点，为了提取边缘像素点的特征）
            ), 
            nn.ReLU(), # 激活函数
            nn.MaxPool2d(kernel_size=2), # 池化：往下一层筛选你重要的部分，比如筛选这个区域的最大值，作为这个区域的特征，这里是导致图片变小的原因
        )
        ################################################
        # 第一层，输入图片的变化：
        # (1, 28, 28) -> nn.Conv2d() -> (16, 28, 28)
        # -> nn.ReLU() -> (16, 28, 28)
        # -> nn.MaxPool2d -> (16, 14, 14)
        #################################################
        # 第二层神经网络
        self.conv2 = nn.Sequential(  # -> (16, 14, 14)
            nn.Conv2d(16, 32, 5, 1, 2),  # -> (32, 14, 14)
            nn.ReLU(),  # -> (32, 14, 14)
            nn.MaxPool2d(2)  # -> (32, 7, 7)
        )
        # 输出层：全连接层
        self.output = nn.Linear(32 * 7 * 7, 10)
    
    # 前向计算
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x) # (batch, 32, 7, 7)
        x = x.view(x.size(0), -1)  # (batch, 32 * 7 * 7) 这里是把图像展平，全连接，相当于reshape
        x = self.output(x)
        return x

源码如下：

# coding=utf-8
####################################################
# 卷积神经网络：
#     手写数字识别
#     搭建卷积神经网络：CNN

# 查看显卡情况：（打开CMD）
#     cd C:\Program Files\NVIDIA Corporation\NVSMI
#     nvidia-smi.exe
####################################################

import os
import sys

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

# 全局变量 ###################################################################
use_cuda = torch.cuda.is_available() # GPU
if use_cuda: # 打印GPU情况
    print(torch.cuda.get_device_name(0), use_cuda)

# 随机种子
manualSeed = 2020
torch.manual_seed(manualSeed)
if use_cuda:
    torch.cuda.manual_seed_all(manualSeed)

# 超参数
EPOCH = 2
BATCH_SIZE = 100
LR = 0.001

save_path = os.path.join(sys.path[0], "cnn.pth") # 保存模型参数的文件

# 显示图片
def display_image(image, label):
    plt.imshow(image, cmap="gray") # 输入一个矩阵
    plt.title('%i' % label.item())
    plt.show()

# 数据分布曲线，用来查看loss的变化
def polt_curve(data): 
    # 数据分布曲线
    fig = plt.figure()
    plt.plot(range(len(data)), data, color = 'blue')
    plt.legend(['value'], loc = 'upper right')
    plt.xlabel('step')
    plt.ylabel('value')
    plt.show()

# 搭建神经网络 ###############################################################
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        # 第一层：输入层
        self.conv1 = nn.Sequential(  # conv1是一个卷积层（一般包括：卷积层，激励函数，池化层），Sequential是其内部的层包裹起来（是一个序贯模型）
            nn.Conv2d(  # 卷积层（2d:二维卷积）：是一个3维过滤器filter，有长宽高，高度指filter的个数（用于提取特征属性）
                in_channels=1, # 输入数据的高度（是上一层的高度，也称卷积核的通道数），图片RGB为3层
                out_channels=16, # 输出的个数（也是filter的个数，是卷积核的个数），是自定义的，表示同时有16个filter在收集信息，这里是使图片变厚的原因
                kernel_size=5, # 表示卷积核的长和宽都是5个像素点
                stride=1,  # 步长：每次卷积核移动的像素点个数(表示每隔多少步跳一下，扫描下一块区域)
                padding=2,  # padding=(kernel_size-stride)/2.表示在图像的边缘填充2个像素点（补全边缘像素点，为了提取边缘像素点的特征）
            ), 
            nn.ReLU(), # 激活函数
            nn.MaxPool2d(kernel_size=2), # 池化：往下一层筛选你重要的部分，比如筛选这个区域的最大值，作为这个区域的特征，这里是导致图片变小的原因
        )
        ################################################
        # 第一层，输入图片的变化：
        # (1, 28, 28) -> nn.Conv2d() -> (16, 28, 28)
        # -> nn.ReLU() -> (16, 28, 28)
        # -> nn.MaxPool2d -> (16, 14, 14)
        #################################################
        # 第二层神经网络
        self.conv2 = nn.Sequential(  # -> (16, 14, 14)
            nn.Conv2d(16, 32, 5, 1, 2),  # -> (32, 14, 14)
            nn.ReLU(),  # -> (32, 14, 14)
            nn.MaxPool2d(2)  # -> (32, 7, 7)
        )
        # 输出层：全连接层
        self.output = nn.Linear(32 * 7 * 7, 10)
    
    # 前向计算
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x) # (batch, 32, 7, 7)
        x = x.view(x.size(0), -1)  # (batch, 32 * 7 * 7) 这里是把图像展平，全连接，相当于reshape
        x = self.output(x)
        return x

# 加载数据集 #################################################################
def data_set(data_path): # 传入数据保存的路径
    # 下载数据集
    train_set = torchvision.datasets.MNIST( #训练集
        root=data_path, # 保存的目录
        train=True, # 是否用于训练
        transform=torchvision.transforms.ToTensor(), # 数据增强、转化，将原始数据改变成什么样的形式，这里主要是将图片的数据点转化成tensor, 同时将像素点的值标准化，压缩到0到1之间
        download=True # 如果当前目录下有数据集，就不会下载
    )

    train_loader = Data.DataLoader(
        dataset=train_set, # 加载数据
        batch_size=BATCH_SIZE, # 每次加载数据的数量
        shuffle=True, # 是否打乱
        num_workers=0 # 工作线程
    )

    test_set = torchvision.datasets.MNIST( # 测试集
        root=data_path,
        train=False,
        transform=torchvision.transforms.ToTensor(),
        download=True
    )

    test_loader = Data.DataLoader(
        dataset=test_set,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=0
    )

    print(f'训练数据集：{train_set.train_data.size()}') # 打印训练数据量
    # 训练数据集：torch.Size([60000, 28, 28])
    print(f'测试数据集：{test_set.test_data.size()}')
    # 测试数据集：torch.Size([10000, 28, 28])

    # # 显示一张train_set的一个数据
    # display_image(train_set.train_data[0].numpy(), train_set.train_labels[0])

    # # 显示一张test_loader的一个数据
    # x, y = next(iter(test_loader)) # 返回的是一个批次的数据
    # ##
    # # x 是一批图片 shape is [batch, 1, 28, 28]
    # # x[0] 是这批图片的第一张图片 shape is [1, 28, 28] 只有一个通道的灰色照片
    # # x[0][0] shape is [28, 28]
    # # y[0]是一个tensor
    # display_image(x[0][0], y[0])

    return train_loader, test_loader

# 训练网络模型 ###############################################################
def train_CNN(train_loader): # 传入训练集
    # 训练网络
    cnn = CNN()
    if use_cuda:
        cnn = cnn.cuda()
    # print(cnn)
    ''' 打印结果
    CNN(
    (conv1): Sequential(
        (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
        (1): ReLU()
        (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (conv2): Sequential(
        (0): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
        (1): ReLU()
        (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (output): Linear(in_features=1568, out_features=10, bias=True)
    )
    '''
    # 分割线 ######################################################################
    # 定义优化器和损失函数
    optimization = torch.optim.Adam(cnn.parameters(), LR)
    loss_func = nn.CrossEntropyLoss() # 交叉熵

    loss_data = []
    # 训练网络
    for epoch in range(EPOCH):
        for step, (x, y) in enumerate(train_loader):
            if use_cuda:
                x, y = x.cuda(), y.cuda()

            out = cnn(x)
            loss = loss_func(out, y)
            # loss_data.append(loss.item())

            optimization.zero_grad()
            loss.backward()
            optimization.step()

            # 每50步打印一下结果
            step += 1
            if step % 50 == 0:
                loss_data.append(loss.item())
                print(f'Epoch:{epoch + 1} Step:{step} Train loss:{loss.item()}')
                # Epoch:10 Step:600 Train loss:0.015473434701561928

    # 打印误差曲线
    polt_curve(loss_data)

    # 保存模型参数
    torch.save(cnn.state_dict(), save_path)
    print(f'网络参数保存成功：{save_path}')
    # 网络参数保存成功：d:\YDDUONG\Ydduong\VSPython\A-GCN-N\test-dir\cnn.pth

# 提取和测试网络 ##############################################################
def test_CNN(test_loader): # 传入测试集
    # 模型提取和测试
    pre_cnn = CNN()
    pre_cnn.load_state_dict(torch.load(save_path))
    if use_cuda:
        pre_cnn = pre_cnn.cuda()

    # 测试通过率
    acc_num = 0 # 预测正确的数目
    for step, (x, y) in enumerate(test_loader):
        if use_cuda:
            x, y = x.cuda(), y.cuda()
        out = pre_cnn(x)
        pred = out.argmax(dim=1) # 预测值
        acc_num += pred.eq(y).sum().item() # 比较、得到正确的

        step += 1
        if step % 10 == 0:
            print(f'Step:{step} Pred miss:{step * 100 - acc_num}')
            # Step:100 Pred miss:88

    print(f'准确率为：{acc_num / len(test_loader.dataset)}')
    # 准确率为：0.9912

# 主流程 #####################################################################
def main():
    # 1. 加载数据
    data_path = sys.path[0] # 数据保存的目录
    train_loader, test_loader = data_set(data_path)

    # 2.训练网络
    train_CNN(train_loader)

    # 3.测试网络
    test_CNN(test_loader)

if __name__ == "__main__":
    main()

参考莫烦pytorch教程：https://space.bilibili.com