Pytorch实现MNIST数据集的训练

随手写了一个用Pytorch实现MNIST数据集的训练模型

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

#下载数据集
from torch.utils.tensorboard import SummaryWriter

train_dataset = torchvision.datasets.MNIST("mnist_data", train=True, transform=torchvision.transforms.ToTensor(),
                                           download=True)
test_dataset = torchvision.datasets.MNIST("mnist_data", train=False, transform=torchvision.transforms.ToTensor(),
                                          download=True)
train_dataset_size = len(train_dataset)
test_dataset_size = len(test_dataset)
print(f"训练数据集的长度为：{train_dataset_size}")
print(f"测试数据集的长度为：{test_dataset_size}")

#载入数据集
train_dataloader = DataLoader(train_dataset, batch_size=64)
test_dataloader = DataLoader(test_dataset, batch_size=64)

#查看图形形状
# for data in test_dataloader:
#     imgs, targets = data
#     print(imgs.shape)

#建立网络
class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.model1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(stride=2, kernel_size=2)
        )

        self.model2 = nn.Sequential(
            nn.Linear(14*14*128, 1024),
            nn.ReLU(),
            nn.Dropout2d(p=0.5),
            nn.Linear(1024, 10)
        )

    def forward(self, x):
        x = self.model1(x)
        x = x.view(-1, 14*14*128)
        x = self.model2(x)
        return x

network = Network()
if torch.cuda.is_available():
    network.cuda()

#建立损失函数
loss_fn = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_fn.cuda()

#建立优化器
learning_rate = 0.01
optimizer = torch.optim.Adam(network.parameters(), lr=learning_rate)

#设置训练次数
total_train_step = 0

#设置测试次数
total_test_step = 0

#训练的次数
epoch = 10

#添加到tensoboard
writer = SummaryWriter("mnist_train_image")

for i in range(epoch):
    print(f"------------第{i + 1}轮训练-------------")

    network.train()
    for data in train_dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.cuda()
            targets = targets.cuda()
        output = network(imgs)
        loss = loss_fn(output, targets)

        #优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        if total_train_step % 100 == 0:
            print(f"训练次数：{total_train_step}  loss：{loss.item()}")
            writer.add_scalar("train_loss", loss.item(), total_train_step)


    #测试步骤开始
    network.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            if torch.cuda.is_available():
                imgs = imgs.cuda()
                targets = targets.cuda()
            output = network(imgs)
            loss = loss_fn(output, targets)
            total_test_loss = total_test_loss + loss
            accuracy = (output.argmax(1) == targets).sum()
            total_accuracy += accuracy

    print(f"整体上的损失函数之和：{total_test_loss}")
    print(f"此轮训练的精度为：{total_accuracy/test_dataset_size}")
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy, total_test_step)
    total_test_step += 1


    #保存参数
    torch.save(network, "network_{}.pth".format(i))
    print("模型已保存")

writer.close()

写的有什么不好的地方还请各位大佬指出，大家一起相互学习！一起相互进步！！
还有就是这个网络可到达精度为98%，可能还有提升的空间。