一、完整代码

1、CNN.py

import torch


# 神经网络定义
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()  # 调用父类的构造方法
        self.conv = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=5, padding=2),  # 1个输入通道，32个输出通道，卷积核大小5*5，padding=2
            torch.nn.BatchNorm2d(32),  # 批标准化
            torch.nn.ReLU(),  # 激活函数,作用：将负数转换为0
            torch.nn.MaxPool2d(2)  # 最大池化
        )

        self.fc = torch.nn.Linear(14 * 14 * 32, 10)  # 全连接层, 14*14*32个输入，10个输出

    def forward(self, x):
        out = self.conv(x)  # 卷积
        out = out.view(out.size(0), -1)  # 展平，变成一维
        out = self.fc(out)  # 全连接
        return out

2、demo_cls.py

import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils

# 1、初始化数据 data
train_data = dataset.MNIST(root="./mnist",  # 数据存放的路径
                           train=True,  # 训练集
                           transform=transforms.ToTensor(),  # 数据类型转换
                           download=True)  # 如果数据不存在，是否下载

test_data = dataset.MNIST(root="./mnist",  # 数据存放的路径
                          train=False,  # 测试集
                          transform=transforms.ToTensor(),  # 数据类型转换
                          download=False)  # 如果数据不存在，是否下载
# batchsize 选择部分数据
train_loader = data_utils.DataLoader(dataset=train_data,  # 数据
                                     batch_size=64,  # 每次取64个数据
                                     shuffle=True)  # 是否打乱数据

test_loader = data_utils.DataLoader(dataset=test_data,  # 数据
                                    batch_size=64,  # 每次取64个数据
                                    shuffle=True)  # 是否打乱数据

# 2、定义一个神经网络 net
from CNN import CNN

cnn = CNN()
cnn = cnn.cuda()

# 3. loss
loss_func = torch.nn.CrossEntropyLoss()  # 交叉熵损失函数

# 4. optimizer
optimizer = torch.optim.Adam(cnn.parameters(), lr=0.01)  # Adam优化器

for epoch in range(10):
    # 5. training
    for i, (images, labels) in enumerate(train_loader):
        images = images.cuda()  # 将数据放到GPU上
        labels = labels.cuda()  # 将数据放到GPU上

        outputs = cnn(images)  # 前向传播
        loss = loss_func(outputs, labels)  # 计算损失

        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

    print("epoch is {}, ite is {}/{}, loss is {}".format(epoch + 1, i, len(train_data) // 64, loss.item()))

    # 6. eval/test
    loss_test = 0
    accuracy = 0
    for i, (images, labels) in enumerate(test_loader):
        images = images.cuda()  # 将数据放到GPU上
        labels = labels.cuda()  # 将数据放到GPU上

        outputs = cnn(images)  # 前向传播
        loss_test += loss_func(outputs, labels)  # 计算损失
        _, pred = outputs.max(1)  # 找到概率最大的下标
        accuracy += (pred == labels).sum().item()  # 计算准确率

    accuracy = accuracy / len(test_data)  # len(test_data)=10000, accuracy=0.98
    loss_test = loss_test / (len(test_data) // 64)  # (len(test_data) // 64) = 156，loss_test=0.0555
    print("epoch is {}, accuracy is {}, loss test is {}".format(epoch + 1, accuracy, loss_test.item()))

# 7. save model
torch.save(cnn, "model/mnist_model.pkl")

3、demo_cls_inference.py

import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as data_utils
from CNN import CNN

test_data = dataset.MNIST(root="./mnist",  # 数据存放的路径
                          train=False,  # 测试集
                          transform=transforms.ToTensor(),  # 数据类型转换
                          download=False)  # 如果数据不存在，是否下载
# batchsize 选择部分数据
test_loader = data_utils.DataLoader(dataset=test_data,  # 数据
                                    batch_size=64,  # 每次取64个数据
                                    shuffle=True)  # 是否打乱数据

cnn = torch.load("model/mnist_model.pkl")
cnn = cnn.cuda()

accuracy = 0

import cv2  # pip install opencv-python

for i, (images, labels) in enumerate(test_loader):
    images = images.cuda()  # 将数据放到GPU上
    labels = labels.cuda()  # 将数据放到GPU上

    outputs = cnn(images)  # 前向传播 (64, 10)
    _, pred = outputs.max(1)  # 找到概率最大的下标 (64,)
    accuracy += (pred == labels).sum().item()  # 计算准确率

    images = images.cpu().numpy()  # 将数据放到CPU上, (64, 1, 28, 28)， 64个样本，1个通道，28*28的图片
    labels = labels.cpu().numpy()  # (64, 1, 28, 28)， labels是一个一维数组，记录了每个样本的真实标签
    pred = pred.cpu().numpy()  # (64,)  pred是一个一维数组，记录了每个样本的预测结果

    # batch_size * 1 * 28 * 28
    for idx in range(images.shape[0]):  # images.shape[0] = 64
        im_data = images[idx]  # (1, 28, 28)
        im_label = labels[idx]  # 标签数值：如2
        im_pred = pred[idx]  # 预测数值：如2
        im_data = im_data.transpose(1, 2, 0)  # 28 * 28 * 1

        print("label:", im_label)
        print("pred:", im_pred)
        cv2.imshow("imdata", im_data)  # 显示图片
        cv2.waitKey(0)  # 等待按键

accuracy = accuracy / len(test_data)  # len(test_data)=10000, accuracy=0.98
print("accuracy:", accuracy)

二、效果图

运行 demo_cls_inference.py

目前对cnn这一块的逻辑还不太理解，因为刚入门AI这个领域才几天

准备先系统学下 “计算机系统视觉入门”

参考资料：https://coding.imooc.com/class/chapter/440.html