P1.Pytorch实现mnist手写数字识别

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍖 原作者：K同学啊
  我的环境
  语言环境：python 3.7.12
  编译器：VS code
  深度学习环境：tensorflow 2.7.0、torch1.12.1
  数据：本地数据集46-data

一、代码

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision

# 设置硬件设备，如果有GPU则使用，没有则使用cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 导入数据
train_ds = torchvision.datasets.MNIST('data',
                                      train=True,
                                      transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
                                      download=True)

test_ds  = torchvision.datasets.MNIST('data',
                                      train=False,
                                      transform=torchvision.transforms.ToTensor(), # 将数据类型转化为Tensor
                                      download=True)

# 加载数据
batch_size = 32

train_dl = torch.utils.data.DataLoader(train_ds,
                                       batch_size=batch_size,
                                       shuffle=True)

test_dl  = torch.utils.data.DataLoader(test_ds,
                                       batch_size=batch_size)

# 取一个批次查看数据格式
# 数据的shape为：[batch_size, channel, height, weight]
# 其中batch_size为自己设定，channel，height和weight分别是图片的通道数，高度和宽度。
imgs, labels = next(iter(train_dl))
print(imgs.shape)

# 数据可视化
import numpy as np

# 指定图片大小，图像大小为20宽、5高的绘图(单位为英寸inch)
plt.figure(figsize=(20, 5))
for i, imgs in enumerate(imgs[:20]):
    # 维度缩减
    npimg = np.squeeze(imgs.numpy())
    # 将整个figure分成2行10列，绘制第i+1个子图。
    plt.subplot(2, 10, i + 1)
    plt.imshow(npimg, cmap=plt.cm.binary)
    plt.axis('off')

# plt.show()  如果你使用的是Pycharm编译器，请加上这行代码

# 构建CNN网络
import torch.nn.functional as F

num_classes = 10  # 图片的类别数


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        # 特征提取网络
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)  # 第一层卷积,卷积核大小为3*3
        self.pool1 = nn.MaxPool2d(2)  # 设置池化层，池化核大小为2*2
        self.conv2 = nn.Conv2d(32, 6