PyTorch入门——实现MNIST分类

import torch
from matplotlib import pyplot as plt
from torch import nn, optim
# from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from tqdm import tqdm
from matplotlib.ticker import MaxNLocator

# 超参数
batch_size = 256  # 批大小
learning_rate = 0.0001  # 学习率
epochs = 20  # 迭代次数
channels = 1  # 图像通道大小

# 数据集下载和预处理
transform = transforms.Compose([transforms.ToTensor(),  # 将图片转换成PyTorch中处理的对象Tensor,并且进行标准化0-1
                                transforms.Normalize([0.5], [0.5])])  # 归一化处理
path = './data/'  # 数据集下载后保存的目录
# 下载训练集和测试集
trainData = datasets.MNIST(path, train=True, transform=transform, download=True)
testData = datasets.MNIST(path, train=False, transform=transform)
# 处理成data loader
trainDataLoader = torch.utils.data.DataLoader(dataset=trainData, batch_size=batch_size, shuffle=True)  # 批量读取并打乱
testDataLoader = torch.utils.data.DataLoader(dataset=testData, batch_size=batch_size)


# 开始构建cnn模型
class cnn(torch.nn.Module):
    def __init__(self):
        super(cnn, self).__init__()
        self.model = torch.nn.Sequential(
            # The size of the picture is 28*28
            torch.nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),

            # The size of the picture is 14*14
            torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            #
            # The size of the picture is 7*7
            torch.nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            # torch.nn.MaxPool2d(kernel_size=2, stride=2),

            # The size of the picture is 7*7
            torch.nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1),
            torch.nn.ReLU(),
            # torch.nn.MaxPool2d(kernel_size=2, stride=2),

            torch.nn.Flatten(),
            torch.nn.Linear(in_features=7 * 7 * 256, out_features=512),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.2),  # 抑制过拟合 随机丢掉一些节点
            torch.nn.Linear(in_features=512, out_features=10),
            # torch.nn.Softmax(dim=1) # pytorch的交叉熵函数其实是softmax-log-NLL 所以这里的输出就不需要再softmax了
        )

    def forward(self, input):
        output = self.model(input)
        return output


# 选择模型
model = cnn()
# GPU可用时转到cuda上执行
if torch.cuda.is_available():
    model = model.cuda()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 选用交叉熵函数作为损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# optimizer = optim.Adam(model.parameters())

# 训练模型并存储训练时的指标
epoch = 1
history = {'Train Loss': [],
           'Test Loss': [],
           'Train Acc': [],
           'Test Acc': []}
for epoch in range(1, epochs+1):
    processBar = tqdm(trainDataLoader, unit='step')
    model.train(True)
    train_loss, train_correct = 0, 0
    for step, (train_imgs, labels) in enumerate(processBar):

        if torch.cuda.is_available():  # GPU可用
            train_imgs = train_imgs.cuda()
            labels = labels.cuda()
        model.zero_grad()  # 梯度清零
        outputs = model(train_imgs)  # 输入训练集
        loss = criterion(outputs, labels)  # 计算损失函数
        predictions = torch.argmax(outputs, dim=1)  # 得到预测值
        correct = torch.sum(predictions == labels)
        accuracy = correct / labels.shape[0]  # 计算这一批次的正确率
        loss.backward()  # 反向传播
        optimizer.step()  # 更新优化器参数
        processBar.set_description("[%d/%d] Loss: %.4f, Acc: %.4f" %  # 可视化训练进度条设置
                                   (epoch, epochs, loss.item(), accuracy.item()))

        # 记录下训练的指标
        train_loss = train_loss + loss
        train_correct = train_correct + correct

        # 当所有训练数据都进行了一次训练后，在验证集进行验证
        if step == len(processBar) - 1:
            tst_correct, totalLoss = 0, 0
            model.train(False)  # 开始测试
            model.eval()  # 固定模型的参数并在测试阶段不计算梯度
            with torch.no_grad():
                for test_imgs, test_labels in testDataLoader:
                    if torch.cuda.is_available():
                        test_imgs = test_imgs.cuda()
                        test_labels = test_labels.cuda()
                    tst_outputs = model(test_imgs)
                    tst_loss = criterion(tst_outputs, test_labels)
                    predictions = torch.argmax(tst_outputs, dim=1)

                    totalLoss += tst_loss
                    tst_correct += torch.sum(predictions == test_labels)

                train_accuracy = train_correct / len(trainDataLoader.dataset)
                train_loss = train_loss / len(trainDataLoader)  # 累加loss后除以步数即为平均loss值

                test_accuracy = tst_correct / len(testDataLoader.dataset)  # 累加正确数除以样本数即为验证集正确率
                test_loss = totalLoss / len(testDataLoader)  # 累加loss后除以步数即为平均loss值

                history['Train Loss'].append(train_loss.item())  # 记录loss和acc
                history['Train Acc'].append(train_accuracy.item())
                history['Test Loss'].append(test_loss.item())
                history['Test Acc'].append(test_accuracy.item())

                processBar.set_description("[%d/%d] Loss: %.4f, Acc: %.4f, Test Loss: %.4f, Test Acc: %.4f" %
                                           (epoch, epochs, train_loss.item(), train_accuracy.item(), test_loss.item(),
                                            test_accuracy.item()))
    processBar.close()

# 对测试Loss进行可视化
plt.plot(history['Test Loss'], color='red', label='Test Loss')
plt.plot(history['Train Loss'], label='Train Loss')
plt.legend(loc='best')
plt.grid(True)
plt.xlabel('Epoch')
plt.xlim([0, epoch])
plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
plt.ylabel('Loss')
plt.title('Train and Test LOSS')
plt.legend(loc='upper right')
plt.savefig('LOSS')
plt.show()

# 对测试准确率进行可视化
plt.plot(history['Test Acc'], color='red', label='Test Acc')
plt.plot(history['Train Acc'], label='Train Acc')
plt.legend(loc='best')
plt.grid(True)
plt.xlabel('Epoch')
plt.xlim([0, epoch])
plt.gca().xaxis.set_major_locator(MaxNLocator(integer=True))
plt.ylabel('Accuracy')
plt.title('Train and Test ACC')
plt.legend(loc='lower right')
plt.savefig('ACC')
plt.show()

torch.save(model, './model.pth')