#AI挑战营第一站#基于Pytorch的卷积神经网络识别MNIST数据集-小白从零开始篇

爱吃鱼的加菲猫

#AI挑战营第一站#基于Pytorch的卷积神经网络识别MNIST数据集-小白从零开始篇 [复制链接]

作为一个AI领域纯纯的小白，想要用Pytorch基于MNIST数据集训练个卷积网络模型，本帖子也作为我的学习log吧，高手们轻点拍。

• 先上本人电脑配置：

CPU：Intel(R) Core(TM) i5-4460  CPU @ 3.20GHz   3.20 GHz 四核四线程

GPU：NVIDIA GeForce GTX 750 Ti（未装CUDA）

RAM：8G

Windows 10 专业版

Python 3.12.1

• 搭建CUDA环境：

cmd命令行内输入nvidia-smi.exe，查看本机显卡驱动版本及可用的CUDA版本。

 

这里要注意，Pytorch对CUDA版本有要求，不能乱装，所以得到官网看看匹配的版本。

https://pytorch.org/get-started/locally/

 

这里可以看到，最新稳定版本的PyTorch能支持到CUDA12.1版本。

把我参考的安装教程链接贴上来，具体大家可以照着操作。

CUDA官网下载链接：https://developer.nvidia.com/cuda-toolkit-archive

 

根据自己系统进行选择，下载安装包。

  下载完参照教程双击安装即可，接下来确认下环境变量。

 

cmd命令行里输入nvcc -version命令查看软件版本。

 

• 安装Anaconda

参考安装教程：https://blog.csdn.net/MCYZSF/article/details/116525159

anaconda是一个集成的工具软件，管理它安装Pytorch会很方便https://www.anaconda.com/

再来个GPU版本使用清华源的安装教程。https://cloud.tencent.com/developer/article/2405758

教程二https://www.zhihu.com/tardis/zm/art/681417713?source_id=1005

开始菜单栏找到Anaconda Prompt，我前面创建了虚拟环境名字叫pytorch，输入python，输入activate pytorch切换环境，输入import torch后回车，再输入torch.cuda.is_available()回车，

如果返回True，则证明环境已经配置好。

 

• 模型训练

接下来就可以正式开始模型训练工作了，这里我参考的实现是下面的教程，非常详细。

https://blog.csdn.net/qq_45588019/article/details/120935828

教程二：https://zhuanlan.zhihu.com/p/137571225

1. 导入相关包文件

import torch
import torchvision
from torch.utils.data import DataLoader

2. 引入MNIST数据集

定义超参数

n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
torch.manual_seed(random_seed)

下载获取数据集（训练集及测试集），并用Dataloader 包起来，即载入数据集

train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size_train, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('./data/', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize(
                                       (0.1307,), (0.3081,))
                               ])),
    batch_size=batch_size_test, shuffle=True)

运行上面的程序后，会自动将数据集下载到目录下的data文件夹。

3. 展示MNIST数据集（如果要直接训练，这步可以忽略）

首先看看测试数据有啥。

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
print(example_targets)
print(example_data.shape)

安装相关包文件，在Anaconda Prompt下执行以下命令：

conda install matplotlib

举例展示6幅图，包含图片内容和标签。

fig = plt.figure()
for i in range(12):
    plt.subplot(3, 4, i+1)
    plt.tight_layout()
    plt.imshow(train_dataset.train_data[i], cmap='gray', interpolation='none')
    plt.title("Labels: {}".format(train_dataset.train_labels[i]))
    plt.xticks([])
    plt.yticks([])
plt.show()

 

 

3. 构建模型

此处使用两个2d卷积层，然后是两个全连接(或线性)层。作为激活函数，我们将选择整流线性单元(简称ReLUs)，作为正则化的手段，我们将使用两个dropout层。

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

 

4. 实例化模型

network = Net()

5. 初始化网络及优化器

criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)  # lr学习率，momentum冲量

 

6. 定义训练及测试轮

• Step1：前馈（forward propagation）

• Step2：反馈（backward propagation）

• Step3：更新（update）

训练轮代码：

 

def train(epoch):
    running_loss = 0.0  # 这整个epoch的loss清零
    running_total = 0
    running_correct = 0
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data
        optimizer.zero_grad()

        # forward + backward + update
        outputs = model(inputs)
        loss = criterion(outputs, target)

        loss.backward()
        optimizer.step()

        # 把运行中的loss累加起来，为了下面300次一除
        running_loss += loss.item()
        # 把运行中的准确率acc算出来
        _, predicted = torch.max(outputs.data, dim=1)
        running_total += inputs.shape[0]
        running_correct += (predicted == target).sum().item()

        if batch_idx % 300 == 299:  # 不想要每一次都出loss，浪费时间，选择每300次出一个平均损失,和准确率
            print('[%d, %5d]: loss: %.3f , acc: %.2f %%'
                  % (epoch + 1, batch_idx + 1, running_loss / 300, 100 * running_correct / running_total))
            running_loss = 0.0  # 这小批300的loss清零
            running_total = 0
            running_correct = 0  # 这小批300的acc清零

完整代码：

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from torchsummary import summary
import time

    
# 创建神经网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.output_layer = nn.Linear(32*7*7, 10)
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.reshape(x.size(0), -1)
        output = self.output_layer(x)
        return output

# 超参数
EPOCH = 2
BATCH_SIZE = 100
LR = 0.001
DOWNLOAD = False # 若已经下载mnist数据集则设为False

# 下载mnist数据
train_data = datasets.MNIST(
    root='./data', # 保存路径
    train=True, # True表示训练集，False表示测试集
    transform=transforms.ToTensor(), # 将0~255压缩为0~1
    download=DOWNLOAD
)

# 旧的写法
print(train_data.train_data.size())
print(train_data.train_labels.size())

# 新的写法
print(train_data.data.size())
print(train_data.targets.size())

# 打印部分数据集的图片
for i in range(2):
    print(train_data.targets[i].item())
    plt.imshow(train_data.data[i].numpy(), cmap='gray')
    plt.show()
    
# DataLoader
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=2
)

# 如果train_data下载好后，test_data也就下载好了
test_data = datasets.MNIST(
    root='./data',
    train=False
)

print(test_data.data.size())
print(test_data.targets.size())

# 新建网络
cnn = CNN()
# 将神经网络移到GPU上
cnn.cuda()
print(cnn)

# 查看网络的结构
model = CNN()
if torch.cuda.is_available():
    model.cuda()
summary(model, input_size=(1,28,28))

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

# 损失函数
loss_func = nn.CrossEntropyLoss()

# 为了节约时间，只使用测试集的前2000个数据
test_x = Variable(
    torch.unsqueeze(test_data.data, dim=1),
    volatile=True
).type(torch.FloatTensor)[:2000]/255 # 将将0~255压缩为0~1

test_y = test_data.targets[:2000]

# # 使用所有的测试集
# test_x = Variable(
#     torch.unsqueeze(test_data.test_data, dim=1),
#     volatile=True
# ).type(torch.FloatTensor)/255 # 将将0~255压缩为0~1

# test_y = test_data.test_labels

# 将测试数据移到GPU上
test_x = test_x.cuda()
test_y = test_y.cuda()

# 开始计时
start = time.time()

# 训练神经网络
for epoch in range(EPOCH):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        # 将训练数据移到GPU上
        batch_x = batch_x.cuda()
        batch_y = batch_y.cuda()
        output = cnn(batch_x)
        loss = loss_func(output, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 每隔50步输出一次信息
        if step%50 == 0:
            test_output = cnn(test_x)
            # 将预测结果移到GPU上
            predict_y = torch.max(test_output, 1)[1].cuda().data.squeeze()
            accuracy = (predict_y == test_y).sum().item() / test_y.size(0)
            print('Epoch', epoch, '|', 'Step', step, '|', 'Loss', loss.data.item(), '|', 'Test Accuracy', accuracy)

# 结束计时
end = time.time()

# 训练耗时
print('Time cost:', end - start, 's')

# 预测
test_output = cnn(test_x[:100])
# 为了将CUDA tensor转化为numpy，需要将数据移回CPU上
# 否则会报错：TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
predict_y = torch.max(test_output, 1)[1].cpu().data.numpy().squeeze()
real_y = test_y[:100].cpu().numpy()
print(predict_y)
print(real_y)

# 打印预测和实际结果
for i in range(10):
    print('Predict', predict_y[i])
    print('Real', real_y[i])
    plt.imshow(test_data.data[i].numpy(), cmap='gray')
    plt.show()

训练结果如下：

 

用时100s左右，精度可以到0.97了，足够手写识别。

 

• 保存并转换成ONNX模型

# 保存模型的状态字典
torch.save(model.state_dict(), 'mnist.pth')
model.load_state_dict(torch.load('mnist.pth'))
# 将模型设置为评估模式（如果需要）
model.eval()
#导出为onnx模型
input = torch.randn(1, 28, 28)
torch.onnx.export(model, input, "mnist.onnx", verbose=True)

mnist.pth (115.43 KB, 下载次数: 0)

2024-4-21 16:29 上传

点击文件名下载附件

mnist.onnx (131.38 KB, 下载次数: 0)

2024-4-21 16:29 上传

点击文件名下载附件

爱吃鱼的加菲猫

感觉我的模型有点问题，大小跟大家的相差挺多

nmg

爱吃鱼的加菲猫 发表于 2024-4-21 16:31 感觉我的模型有点问题，大小跟大家的相差挺多

好像差距是挺大的啊，刚随意看了下，他们上传的同格式的，都是M级别大小的

爱吃鱼的加菲猫

nmg 发表于 2024-4-22 10:09 好像差距是挺大的啊，刚随意看了下，他们上传的同格式的，都是M级别大小的

是的，我觉得也有问题，晚上有空再折腾看看

爱吃鱼的加菲猫

前几天参照网友教程，我是Anaconda Prompt终端命令下进行的，在生成模型文件大小有问题，别人都是MB级别，我的是KB，估计是有问题，今晚又折腾了下，参照网友wangerxian帖子https://bbs.eeworld.com.cn/thread-1278516-1-1.html操作，先安装了PyCharm，然后再执行py。

执行的代码直接copy网友knv的，用GPU参与训练的，GPU占有率大概能到45%左右，CPU几乎跑满了。

生成文件如下：

 

文件大小应该是对了，跟knv生成的模型大小基本一致。

接下来参照其他网友的执行看看：

1、新建个工程

 

环境配置参照上面图片，Path to conda那里需要选择你Anaconda安装目录，别搞错了。Environment选择你生成的虚拟python环境，要提前在Prompt里面安装好相关包文件。

2、新建Python文件

①位置工程文件夹上右键，②新建，③选择Python文件

 

输入个文件名

 

把代码拷贝粘贴进来，然后点开左下角终端按钮

 

 

执行python文件

文件上右键，选择运行

 

执行结果：

 

 

pt文件生成了，但是转onnx时候报错，也不知道啥情况，有可能跟软件版本有关。

再来试试其他GPU训练的代码：

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from torchsummary import summary
import time

    
# 创建神经网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.output_layer = nn.Linear(32*7*7, 10)
        
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.reshape(x.size(0), -1)
        output = self.output_layer(x)
        return output

# 超参数
EPOCH = 2
BATCH_SIZE = 100
LR = 0.001
DOWNLOAD = True # 若已经下载mnist数据集则设为False

# 下载mnist数据
train_data = datasets.MNIST(
    root='./data', # 保存路径
    train=True, # True表示训练集，False表示测试集
    transform=transforms.ToTensor(), # 将0~255压缩为0~1
    download=DOWNLOAD
)

# 旧的写法
print(train_data.train_data.size())
print(train_data.train_labels.size())

# 新的写法
print(train_data.data.size())
print(train_data.targets.size())

# 打印部分数据集的图片
for i in range(2):
    print(train_data.targets[i].item())
    plt.imshow(train_data.data[i].numpy(), cmap='gray')
    plt.show()
    
# DataLoader
train_loader = Data.DataLoader(
    dataset=train_data,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=2
)

# 如果train_data下载好后，test_data也就下载好了
test_data = datasets.MNIST(
    root='./data',
    train=False
)

print(test_data.data.size())
print(test_data.targets.size())

# 新建网络
cnn = CNN()
# 将神经网络移到GPU上
cnn.cuda()
print(cnn)

# 查看网络的结构
model = CNN()
if torch.cuda.is_available():
    model.cuda()
summary(model, input_size=(1,28,28))

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

# 损失函数
loss_func = nn.CrossEntropyLoss()

# 为了节约时间，只使用测试集的前2000个数据
test_x = Variable(
    torch.unsqueeze(test_data.data, dim=1),
    volatile=True
).type(torch.FloatTensor)[:2000]/255 # 将将0~255压缩为0~1

test_y = test_data.targets[:2000]

# # 使用所有的测试集
# test_x = Variable(
#     torch.unsqueeze(test_data.test_data, dim=1),
#     volatile=True
# ).type(torch.FloatTensor)/255 # 将将0~255压缩为0~1

# test_y = test_data.test_labels

# 将测试数据移到GPU上
test_x = test_x.cuda()
test_y = test_y.cuda()

# 开始计时
start = time.time()

# 训练神经网络
for epoch in range(EPOCH):
    for step, (batch_x, batch_y) in enumerate(train_loader):
        # 将训练数据移到GPU上
        batch_x = batch_x.cuda()
        batch_y = batch_y.cuda()
        output = cnn(batch_x)
        loss = loss_func(output, batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 每隔50步输出一次信息
        if step%50 == 0:
            test_output = cnn(test_x)
            # 将预测结果移到GPU上
            predict_y = torch.max(test_output, 1)[1].cuda().data.squeeze()
            accuracy = (predict_y == test_y).sum().item() / test_y.size(0)
            print('Epoch', epoch, '|', 'Step', step, '|', 'Loss', loss.data.item(), '|', 'Test Accuracy', accuracy)

# 结束计时
end = time.time()

# 训练耗时
print('Time cost:', end - start, 's')

# 预测
test_output = cnn(test_x[:100])
# 为了将CUDA tensor转化为numpy，需要将数据移回CPU上
# 否则会报错：TypeError: can't convert CUDA tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
predict_y = torch.max(test_output, 1)[1].cpu().data.numpy().squeeze()
real_y = test_y[:100].cpu().numpy()
print(predict_y)
print(real_y)

# 打印预测和实际结果
for i in range(10):
    print('Predict', predict_y[i])
    print('Real', real_y[i])
    plt.imshow(test_data.data[i].numpy(), cmap='gray')
    plt.show()

执行后报错，

 

这里可以看到是没有torchsummary包，通过conda尝试，发现找不到这个包，无奈只能用pip来安装，

pip install torchsummary

 再次执行，注意，如果你已经执行下载过mnist数据集，则39行为false，执行后这里会打开图片样本。

 

关闭图片窗口后会报错，搞不定了，继续换其他代码试试。

  这里采用luyism的代码试试

https://bbs.eeworld.com.cn/thread-1278192-1-1.html

这里提示报错，原因是没安装onnxruntime包，在conda下安装下再次尝试可以生成模型文件，

 

 

这里尝试下qiao--网友的代码，

# 导包
import torch
import torch.nn as nn  # 神经网络
import torch.optim as optim  # 定义优化器
from torchvision import datasets, transforms  # 数据集    transforms完成对数据的处理

# 定义超参数
input_size = 28 * 28  # 输入大小
hidden_size = 512  # 隐藏层大小
num_classes = 10  # 输出大小（类别数）
batch_size = 100  # 批大小
learning_rate = 0.001  # 学习率
num_epochs = 10  # 训练轮数

# 加载 MNIST 数据集
train_dataset = datasets.MNIST(root='../data/mnist', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='../data/mnist', train=False, transform=transforms.ToTensor(), download=True)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  # 一批数据为100个
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)


# 定义 MLP 网络
class MLP(nn.Module):
    # 初始化方法
    # input_size 输入数据的维度
    # hidden_size 隐藏层的大小
    # num_classes 输出分类的数量
    def __init__(self, input_size, hidden_size, num_classes):
        # 调用父类的初始化方法
        super(MLP, self).__init__()
        # 定义第1个全连接层
        self.fc1 = nn.Linear(input_size, hidden_size)
        # 定义ReLu激活函数
        self.relu = nn.ReLU()
        # 定义第2个全连接层
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        # 定义第3个全连接层
        self.fc3 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # 将输入张量展平为向量
        x = x.view(x.size(0), -1)
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out


# 实例化 MLP 网络
model = MLP(input_size, hidden_size, num_classes)

# 现在我们已经定义了 MLP 网络并加载了 MNIST 数据集，接下来使用 PyTorch 的自动求导功能和优化器进行训练。首先，定义损失函数和优化器；然后迭代训练数据并使用优化器更新网络参数。

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# CrossEntropyLoss = Softmax + log + nllloss
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# optimizer = optim.SGD(model.parameters(),0.2)

# 训练网络
# 外层for循环控制训练的次数
# 内层for循环控制从DataLoader中循环读取数据
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, 28 * 28)  # 将images转换成向量

        outputs = model(images)  # 将数据送到网络中

        loss = criterion(outputs, labels)  # 计算损失

        optimizer.zero_grad()  # 首先将梯度清零
        loss.backward()  # 反向传播

        optimizer.step()  # 更新参数

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}],Step[{i + 1}/{len(train_loader)}],Loss:{loss.item():.4f}')

# 最后，我们可以在测试数据上评估模型的准确率：
# 测试网络
with torch.no_grad():
    correct = 0
    total = 0
    # 从test_loader中循环读取测试数据
    for images, labels in test_loader:
        # 将images转换成向量
        images = images.reshape(-1, 28 * 28)
        # 将数据传送到网络
        outputs = model(images)
        # 取出最大值对应的索引 即预测值
        _, predicted = torch.max(outputs.data, 1)  # 返回一个元组：第一个为最大值，第二个是最大值的下标
        # 累加labels数量   labels为形状为(batch_size,1)的矩阵，取size(0)就是取出batch_size的大小（该批的大小）
        total += labels.size(0)
        # 预测值与labels值对比 获取预测正确的数量
        correct += (predicted == labels).sum().item()
    # 打印最终的准确率
    print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

# 保存模型
# 保存模型的状态字典
torch.save(model.state_dict(), 'mnist.pth')
model.load_state_dict(torch.load('mnist.pth'))
# 将模型设置为评估模式（如果需要）
model.eval()
#导出为onnx模型
input = torch.randn(1, 28, 28)
torch.onnx.export(model, input, "mnist.onnx", verbose=True)

执行结果如下;

 

我发现大家生成的模型文件差别挺大的，这里模型准确率可以到98%了，10轮训练。

qiao.7z (2.36 MB, 下载次数: 1)

2024-4-23 21:25 上传

点击文件名下载附件

LitchiCheng.7z (215.44 KB, 下载次数: 0)

2024-4-23 21:26 上传

点击文件名下载附件

minst.7z (1.49 MB, 下载次数: 0)

2024-4-23 21:26 上传

点击文件名下载附件
knx