手写MNIST数字识别（pytorch）

可使用GPU进行训练
CPU版参考视频
参考代码（CPU）
参考代码（GPU）

import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt
from torch import nn

# 检查CUDA设备是否可用，然后选择设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#神经网络主体
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        #含4个全连接层
        self.fc1=nn.Linear(28*28,64)#输入为28*28图像,第0层设有64个节点
        self.fc2=nn.Linear(64,64)#0层：64个节点、1层：64个节点
        self.fc3=nn.Linear(64,64)#1层：64个节点，2层：64个节点
        self.fc4=nn.Linear(64,10)#2层：64个节点，输出：10个值
    
    def forward(self,x):#参数x为图像输入
        #每层传播中，先进行全连接线性计算，再套上激活函数
        x=nn.functional.relu(self.fc1(x))
        x=nn.functional.relu(self.fc2(x))
        x=nn.functional.relu(self.fc3(x))
        x=nn.functional.log_softmax(self.fc4(x),dim=1)#沿行（某一样本在所有类别）归一化
        return x

def get_data_loader(is_train):
    #数据转化为tensor类型
    to_tensor=transforms.Compose([transforms.ToTensor()])
    
    # 下载MNIST数据集
    # “”表示下载目录，若没指明则表示当前目录
    # is_train表示用于导入训练集/测试集,它是一个布尔值，指定是加载训练集还是测试集。True 表示加载训练集（60,000 张图片）。False 表示加载测试集（10,000 张图片）。
    # transform 数据预处理方法，用于将数据转换成 PyTorch 张量。
    # download=True，如果数据集不存在，会自动下载。
    data_set=MNIST("",is_train,transform=to_tensor,download=True)
    
    # data_set表示使用的数据集
    # batch_size=15表示一个批次包含15张图片
    # shuffle=True
    return DataLoader(data_set,batch_size=15,shuffle=True)
    # return DataLoader(data_set,batch_size=128,shuffle=True)


# 评估神经网络识别正确率
def evaluate(test_data,net):
        #n_correct：用于记录模型预测正确的样本数量。
        #n_total：用于记录测试样本的总数量。
        n_correct=0
        n_total=0
        with torch.no_grad():#禁用梯度计算。在评估模型时，我们不需要计算梯度，因为不会进行反向传播。使用 torch.no_grad() 可以节省内存并提高计算效率。
            # 在测试集中按批次取出数据
            for (x,y) in test_data:
                # 计算神经网络的预测值
                x,y=x.to(device),y.to(device)
                outputs=net.forward(x.view(-1,28*28))
                # 对批次中的每个结果进行比较，累加正确数量。
                for i,output in enumerate(outputs):
                    # argmax:计算数列中数列最大值的序号
                    if torch.argmax(output)==y[i]:
                        n_correct+=1
                    n_total+=1

                # 使用批量操作计算准确率
                # predictions = torch.argmax(outputs, dim=1)
                # n_correct += (predictions == y).sum().item()
                # n_total += y.size(0)
        return n_correct/n_total

def main():
    train_data=get_data_loader(is_train=True)
    test_data=get_data_loader(is_train=False)
    net=Net()
    net.to(device)# 将模型移动到选择的设备

    print("initial accuracy:",evaluate(test_data,net))

    #训练神经网络，pytorch固定写法
    optimizer=torch.optim.Adam(net.parameters(),lr=0.001)
    for epoch in range(2):
        for (x,y) in train_data:
            x,y=x.to(device),y.to(device)
            net.zero_grad()#初始化
            output=net.forward(x.view(-1,28*28))#正向传播
            loss=nn.functional.nll_loss(output,y)#计算误差损失
            loss.backward()#反向误差传播
            optimizer.step()#优化网络参数
        print("epoch",epoch,"accuracy:",evaluate(test_data,net))
    
    # x: 当前批次的输入图像张量（形状通常为 [batch_size, channels, height, width]）。
    # _: 当前批次的目标标签张量（因为标签在这里没用，所以用 _ 忽略它）
    for (n,(x,_)) in enumerate(test_data):
        if n>3:
            break
        x=x.to(device)
        # x[0]: 从当前批次中选择第一张图像，形状为 [1, 28, 28]（假设输入是 MNIST 数据）。
        # .view(-1, 28*28): .view类似numpy中的reshape，这里将选中的图像展平成一个一维向量，形状变为 [1, 784]，以匹配模型 net 的输入要求。
        # net.forward(x[0].view(-1, 28*28)): 将展平后的图像输入到模型中，得到预测结果。输出是一个形状为 [1, 10] 的张量，表示该图像在 10 个类别上的预测分数。
        # torch.argmax(...): 找到输出张量中分数最高的索引，这个索引代表模型预测的类别。
        # predict: 预测的类别索引（整数），即模型认为这张图片属于哪一类（如 0-9 的数字）。
        predict=torch.argmax(net.forward(x[0].view(-1,28*28)))
        
        # plt.figure(n): 创建一个新的 Matplotlib 图像窗口，用于绘制第 n 个图像。
        plt.figure(n)
        # x[0]: 当前批次的第一张图像，形状为 [1, 28, 28]。
        # .view(28, 28): 将张量调整为二维形状 [28, 28]，以便 Matplotlib 显示为图像。
        # .cpu(): 如果 x[0] 在 GPU 上，需要将其移动到 CPU 上，才能使用 Matplotlib 处理。
        #.numpy(): 将 PyTorch 张量转换为 NumPy 数组，这是 Matplotlib 能够处理的格式。
        # cmap='gray': 指定颜色映射为灰度图，适用于 MNIST 图像数据。
        plt.imshow(x[0].view(28,28).cpu().numpy(),cmap='gray')
        # str(int(predict)): 将预测的类别索引转换为字符串形式，显示在标题中。
        plt.title("prediction: "+str(int(predict)))
    plt.show()


if __name__=="__main__":
    main()