基于Lenet-5的MNIST数据集处理

import torch
import torchvision.datasets
from torch import nn, optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

train_data = torchvision.datasets.MNIST("./data", train=True,
                                        transform=torchvision.transforms.ToTensor(), download=True)
test_data = torchvision.datasets.MNIST("./data", train=False,
                                       transform=torchvision.transforms.ToTensor(), download=True)
train_Dataloader = DataLoader(train_data, batch_size=64, shuffle=True)
test_Dataloader = DataLoader(test_data, batch_size=64, shuffle=False)


class Lenet_5(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
        )
        self.layer3 = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
        )
        self.layer4 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU(),
        )
        self.layer5 = nn.Linear(84, 10)

    def forward(self, input):
        output = self.layer1(input)
        output = self.layer2(output)
        # print(output.shape)
        output = output.view(output.size(0), -1)
        output = self.layer3(output)
        # print("layer3", output.shape)
        output = self.layer4(output)
        output = self.layer5(output)
        return output


test_data_size = len(test_data)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
lenet_5 = Lenet_5()


writer = SummaryWriter("log_lenet_5")
lenet_5.to(device)
# loss_fuc = nn.CrossEntropyLoss().to(device)    #多分类问题，选择交叉熵损失函数
loss_fn = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(lenet_5.parameters(), lr=0.001, momentum=0.9)
epoch = 10
train_step = 0
total_test_step = 0
max_accuracy = 0
for i in range(epoch):
    for data in train_Dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.to(device)
            targets = targets.to(device)
        # 梯度清零
        optimizer.zero_grad()

        # 传递损失 + 更新参数
        output = lenet_5(imgs)
        # print("output_shape", output.shape)
        # print("targets_shape", targets.shape)
        loss = loss_fn(output, targets)
        loss.backward()
        optimizer.step()
        if train_step % 100 ==0:
            print("训练 {} 次:".format(train_step+1)+". loss={}:".format(loss))
            writer.add_scalar("train_loss", loss.item(), train_step)
        train_step = train_step + 1
        if train_step == 3:
            writer.add_graph(lenet_5, imgs)
    total_accuracy = 0
    total_test_loss = 0
    accuracy = 0
    for data in test_Dataloader:
        imgs, targets = data
        if torch.cuda.is_available():
            imgs = imgs.to(device)
            targets = targets.to(device)
        outputs = lenet_5(imgs)
        loss = loss_fn(outputs, targets)
        total_test_loss = total_test_loss + loss.item()
        accuracy = (outputs.argmax(1) == targets.reshape(targets.shape, 1)).sum().to(device)
        # print(outputs.argmax(1))
        # print(targets.reshape(targets.shape, 1))
        total_accuracy = total_accuracy + accuracy
    print("整体测试集上的正确率{}".format(total_accuracy / test_data_size))
    writer.add_scalar("test_loss", total_test_loss, total_test_step)
    writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step)
    total_test_step = total_test_step + 1
    if max_accuracy < total_accuracy:
        # 保存模型
        torch.save(lenet_5.state_dict(), 'model.pkl')
        max_accuracy = total_accuracy
writer.close()