#AI挑战营第一站#基于PyTorch，在PC上完成MNIST手写数字识别模型训练

硬核王同学

#AI挑战营第一站#基于PyTorch，在PC上完成MNIST手写数字识别模型训练 [复制链接]

本帖最后由硬核王同学于 2024-4-24 10:20 编辑

注意！网上有很多版本无法实现手写数字模型训练，我这个亲手实验，是可以成功的！

前几天一直在做环境的安装，我想在Ubuntu14上运行，发现python都不支持，后来升级到22可以了，但是在那个环境里没有显卡，所以还是玩不起来。

索性还是在windows里面训练吧，也记录下如何在PC上完成手写数字模型训练。

1.安装环境

1.1安装CUDA

首先查看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

根据系统选择windows、64、11、exe版本，下载比较慢稍微等一下。

直接双击安装，下一步就好了。

重启cmd，查看nvcc --version软件版本

1.2安装Anaconda

登录anaconda的官网下载，anaconda是一个集成的工具软件不需要我们再次下载。

anaconda官网

输入你的邮箱注册

点击立即下载，选择windows版本

直接双击安装，下一步就好了。

windows搜索框里找一下anaconda prompt

打开，出现（base），安装成功。

1.3配置pytorch环境

在Anaconda Prompt下敲命令：conda create -n python python=3.11

此时需要下载导入一些包，输入y。

导入完成，切换虚拟环境名字叫python：conda activate python，再pip list查看当前环境，发现没有PyTorch，那么下面我们就需要安装它。

1.4安装pytorch

进入pytorch官网，选择本机CUDA12.1版本安装

复制下面的命令：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

在原来的python环境中输入命令，开始下载：

下载完成后我们再次输入pip list，查看到已经有torch的存在。

先输入python，然后输入import torch，如果输入后没有任何报错，没有任何显示那就是成功了，然后再输入torch.cuda.is_available()，返回的是True，那便是完成了整个操作。

2.训练模型

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

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

提前安装好相关包文件，在python环境下执行以下命令：

pip install matplotlib

pip install onnx

2.1导入包：

import torch
import torchvision
from torch.utils.data import DataLoader
import numpy as np
from matplotlib import pyplot as plt
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F
import torch.optim as optim

2.2定义超参数：

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)

2.3载入数据集：

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文件夹。

2.4查看MNIST数据集：

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

fig = plt.figure()
for i in range(6):
  plt.subplot(2,3,i+1)
  plt.tight_layout()
  plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
  plt.title("Ground Truth: {}".format(example_targets[i]))
  plt.xticks([])
  plt.yticks([])

这里记得跳出循环后再敲plt.show()

plt.show()

2.5构建网络：

现在让我们开始建立我们的网络。我们将使用两个2d卷积层，然后是两个全连接(或线性)层。作为激活函数，我们将选择整流线性单元(简称ReLUs)，作为正则化的手段，我们将使用两个dropout层。在PyTorch中，构建网络的一个好方法是为我们希望构建的网络创建一个新类。让我们在这里导入一些子模块，以获得更具可读性的代码。

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)
        x = F.log_softmax(x, dim=1)
        return x

具体各部分的含义，在下面详细讲！

广义地说，我们可以想到torch.nn层中包含可训练的参数，而torch.nn.functional就是纯粹的功能性。forward()传递定义了使用给定的层和函数计算输出的方式。为了便于调试，在前向传递中打印出张量是完全可以的。在试验更复杂的模型时，这就派上用场了。请注意，前向传递可以使用成员变量甚至数据本身来确定执行路径——它还可以使用多个参数!

2.6实例化模型：

network = Net()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
network.to(device)

2.7初始化网络及优化器：

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

2.8模型训练：

建立我们的训练循环了。

首先，我们要确保我们的网络处于训练模式。然后，每个epoch对所有训练数据进行一次迭代。加载单独批次由DataLoader处理。

我们需要使用optimizer.zero_grad()手动将梯度设置为零，因为PyTorch在默认情况下会累积梯度。

然后，我们生成网络的输出(前向传递)，并计算输出与真值标签之间的负对数概率损失。

现在，我们收集一组新的梯度，并使用optimizer.step()将其传播回每个网络参数。

有关PyTorch自动渐变系统内部工作方式的详细信息，请参阅autograd的官方文档(强烈推荐)。

我们还将使用一些打印输出来跟踪进度。为了在以后创建一个良好的培训曲线，我们还创建了两个列表来节省培训和测试损失。在x轴上，我们希望显示网络在培训期间看到的培训示例的数量。

train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

在开始训练之前，我们将运行一次测试循环，看看仅使用随机初始化的网络参数可以获得多大的精度/损失。你能猜出我们的准确度是多少吗?

def train(epoch):
    network.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)  # 将数据和标签移动到设备上
        optimizer.zero_grad()
        output = network(data)
        loss = criterion(output, target)  # 使用交叉熵损失函数
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
                (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
            torch.save(network.state_dict(), './model.pth')
            torch.save(optimizer.state_dict(), './optimizer.pth')

train(1)

神经网络模块以及优化器能够使用.state_dict()保存和加载它们的内部状态。这样，如果需要，我们就可以继续从以前保存的状态dict中进行训练——只需调用.load_state_dict(state_dict)。

2.9循环测试：

现在进入循环测试。在这里，我们总结了测试损失，并跟踪正确分类的数字来计算网络的精度。

def test():
  network.eval()
  test_loss = 0
  correct = 0
  with torch.no_grad():
    for data, target in test_loader:
      data, target = data.to(device), target.to(device)
      output = network(data)
      test_loss += F.nll_loss(output, target, reduction='sum').item()
      pred = output.data.max(1, keepdim=True)[1]
      correct += pred.eq(target.data.view_as(pred)).sum()
  test_loss /= len(test_loader.dataset)
  test_losses.append(test_loss)
  print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss, correct, len(test_loader.dataset),100. * correct / len(test_loader.dataset)))

test()

开始多次训练!

我们将在循环遍历n_epochs之前手动添加test()调用，以使用随机初始化的参数来评估我们的模型。

for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()

2.10评估模型的性能：

经过几个阶段的训练，我们已经能够达到测试集97%的准确率!我们开始使用随机初始化的参数。

我们来画一下训练曲线。

import matplotlib.pyplot as plt
fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()

2.11模型导出为pth文件：

torch.save(network.state_dict(), 'mnist_101_model.pth')

2.12导出Onnx文件：

dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(network, dummy_input.to(device), "mnist_101_model.onnx", verbose=False)

3.完整代码：

import torch
import torchvision
from torch.utils.data import DataLoader
import numpy as np
from matplotlib import pyplot as plt
from torchvision import transforms
from torchvision import datasets
import torch.nn.functional as F
import torch.optim as optim

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)

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)

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)
        x = F.log_softmax(x, dim=1)
        return x
        
network = Net()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
network.to(device)
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失
optimizer = torch.optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)  # lr学习率，momentum冲量

train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]

def train(epoch):
    network.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)  # 将数据和标签移动到设备上
        optimizer.zero_grad()
        output = network(data)
        loss = criterion(output, target)  # 使用交叉熵损失函数
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            train_losses.append(loss.item())
            train_counter.append(
                (batch_idx * 64) + ((epoch - 1) * len(train_loader.dataset)))
            torch.save(network.state_dict(), './model.pth')
            torch.save(optimizer.state_dict(), './optimizer.pth')

train(1)

def test():
  network.eval()
  test_loss = 0
  correct = 0
  with torch.no_grad():
    for data, target in test_loader:
      data, target = data.to(device), target.to(device)
      output = network(data)
      test_loss += F.nll_loss(output, target, reduction='sum').item()
      pred = output.data.max(1, keepdim=True)[1]
      correct += pred.eq(target.data.view_as(pred)).sum()
  test_loss /= len(test_loader.dataset)
  test_losses.append(test_loss)
  print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss, correct, len(test_loader.dataset),100. * correct / len(test_loader.dataset)))

test()

for epoch in range(1, n_epochs + 1):
    train(epoch)
    test()
    
import matplotlib.pyplot as plt
fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()

torch.save(network.state_dict(), 'mnist_101_model.pth')
dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(network, dummy_input.to(device), "mnist_101_model.onnx", verbose=False)

wangerxian

用cuda训练，速度应该很快，一套下来用了多久时间？

硬核王同学

wangerxian 发表于 2024-4-24 09:15 用cuda训练，速度应该很快，一套下来用了多久时间？

3轮训练不到一分钟

wangerxian

硬核王同学发表于 2024-4-24 17:20 3轮训练不到一分钟

那挺快的，我用CPU训练6轮，快十分钟了。

爱吃鱼的加菲猫

wangerxian 发表于 2024-4-24 17:22 那挺快的，我用CPU训练6轮，快十分钟了。

我用CPU训练时候，直接占用率拉满，电脑最好啥事都别干，GPU时候还可以正常使用

wangerxian

爱吃鱼的加菲猫发表于 2024-4-25 10:23 我用CPU训练时候，直接占用率拉满，电脑最好啥事都别干，GPU时候还可以正常使用

嗯嗯，就和电脑录屏一样，录屏用CPU，占用率很高的，用GPU占用率就低得多。

freebsder