Python PyTorch4：MNIST案例

MNIST数据集（Mixed National Institute of Standards and Technology database）是美国国家标准与技术研究院收集整理的大型手写数字数据库，包含60,000个示例的训练集以及10,000个示例的测试集。MNIST数据集包含了四个部分: Training set images (包含 60,000 个样本)、Training set labels (包含 60,000 个标签)、Test set images (包含 10,000 个样本)、 Test set labels (包含 10,000 个标签)。

链接：https://pan.baidu.com/s/1rr1H8EKuhtvCsRXr5cb6Lw
提取码：imn8

MNIST的样本

MNIST的每张图片包含28*28个像素，我们把这一个数组展开成一个向量，长度是28*28=784。 如果把数据用矩阵表示，可以把MNIST训练数据变成一个形状为 [60000, 784] 的矩阵，第一个维度数字用来索引图片，第二个维度数字用来索引每张图片中的像素点。图片里的某个像素的强度值介于0-1之间。

MNIST的标签

MNIST数据集的标签是介于0-9的数字，我们要把标签转化为“one-hot vectors”。一个one-hot向量除了某一位数字是1以外，其余维度数字都是0。比如标签0将表示为([1,0,0,0,0,0,0,0,0,0])，标签3将表示为  ([0,0,0,1,0,0,0,0,0,0]) 。因此，可以把MNIST训练集的标签变为 [60000, 10] 的矩阵。

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

# 训练集
train_dataset = datasets.MNIST(root='./',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=True)
# 测试集
test_dataset = datasets.MNIST(root='./',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=True)

# 批次大小
batch_size = 64
# 装载训练集
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True)
# 装载训练集
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         shuffle=True)

'''for i, data in enumerate(train_loader):
    inputs, labels = data
    # torch.Size([64, 1, 28, 28])
    # 每次循环64个数据 通道数1黑白 28*28像素
    print(inputs.shape)
    # torch.Size([64])
    # 图片的标签 1,2,3...
    print(labels.shape)
    break'''


# 定义网络结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 784=28*28 输入28*28种特征，输出10种特征
        self.fc1 = nn.Linear(784, 10)
        # 分类 使用softmax激活函数
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # ([64, 1, 28, 28])->(64,784)
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.softmax(x)
        return x


# 学习率
LR = 0.5
# 定义模型
model = Net()
# 定义代价函数
mse_loss = nn.MSELoss()
# 定义优化器
optimizer = optim.SGD(model.parameters(), LR)


def train():
    for i, data in enumerate(train_loader):
        # 获得一个批次的数据和标签
        inputs, labels = data
        # 获得模型预测结果（64,10）
        out = model(inputs)
        # (64)-(64,1)
        labels = labels.reshape(-1, 1)

        # to one-hot
        # 生成64行10列的矩阵
        # tensor.scatter(dim, index, src)
        # dim:对哪个维度进行one-hot编码，这里是“10”
        # index:要将src中对应的值放到tensor的哪个位置，如第0个值是3则第0行第3位是1其他是0。
        # src:插入index的数值
        one_hot = torch.zeros(inputs.shape[0], 10).scatter(1, labels, 1)
        # 计算loss,mes_loss的两个数据的shape要一致，所以要one-hot
        loss = mse_loss(out, one_hot)
        # 梯度清0
        optimizer.zero_grad()
        # 计算梯度
        loss.backward()
        # 修改权值
        optimizer.step()


def test():
    correct = 0
    for i, data in enumerate(test_loader):
        # 获得一个批次的数据和标签
        inputs, labels = data
        # 获得模型预测结果（64,10）
        out = model(inputs)
        # 获得最大值，以及最大值所在的位置 1表示out的第1个维度（10）
        _, predicted = torch.max(out, 1)
        # 预测正确的数量
        correct += (predicted == labels).sum()
    print("Test acc:{0}".format(correct.item() / len(test_dataset)))


for epoch in range(10):
    print('epoch:', epoch)
    train()
    test()

输出：

epoch: 0
Test acc:0.8893
epoch: 1
Test acc:0.9014
epoch: 2
Test acc:0.9071
epoch: 3
Test acc:0.9107
epoch: 4
Test acc:0.9139
epoch: 5
Test acc:0.9158
epoch: 6
Test acc:0.9174
epoch: 7
Test acc:0.9183
epoch: 8
Test acc:0.9191
epoch: 9
Test acc:0.9191