Pytorch入门：手写数字识别代码逐行解析

先上完整代码：

import torch
from torch import nn
from torch.utils.data import Dataset
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torch.optim as optim

#定义神经网络类Network, 它将继承nn.Module类
class Network(nn.Module):
    # 神经网络中的神经元数量是固定的，所以init不用传入参数
    def __init__(self):
        super().__init__() # 调用了父类的初始化函数
        # layer1是输入层于隐藏层之间的线性层
        self.layer1 = nn.Linear(28 * 28, 256)
        # layer2是隐藏层和输出层之间的线性层
        self.layer2 = nn.Linear(256, 10)
    def forward(self, x):
        #将输入由二维转换为一维，便于输入到线性层中
        x = x.view(-1, 28*28)
        x = self.layer1(x) #输出layer1的输出结果
        x = torch.relu(x) #输出激活后的输出结果
        #计算layer2的结果，并返回
        return self.layer2(x)

device = torch.device("cuda:0")
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())

# 定义数据加载器
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)


# 定义模型、损失函数和优化器

model = Network().to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters())


# 训练模型
for epoch in range(20):
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 20, i+1, len(train_loader), loss.item()))
            
            
            
# 测试模型
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy: {:.2f}%'.format(100 * correct / total))

下面开始逐行解析

1.

from torch import nn

导入pytorch中的nn，该类包含了构建模型所需要的组件

import torch.optim as optim

optim里有各种优化器

2.

#定义神经网络类Network, 它将继承nn.Module类
class Network(nn.Module):
    # 神经网络中的神经元数量是固定的，所以init不用传入参数
    def __init__(self):
        super().__init__() # 调用了父类的初始化函数
        # layer1是输入层于隐藏层之间的线性层
        self.layer1 = nn.Linear(28 * 28, 256)
        # layer2是隐藏层和输出层之间的线性层
        self.layer2 = nn.Linear(256, 10)
    def forward(self, x):
        #将输入由二维转换为一维，便于输入到线性层中
        x = x.view(-1, 28*28)
        x = self.layer1(x) #输出layer1的输出结果
        x = torch.relu(x) #输出激活后的输出结果
        #计算layer2的结果，并返回
        return self.layer2(x)

创建模型类

先在构造方法__init__()中创建搭建模型所需要的各种层。然后在forward()方法中搭建模型并实现前向传播的过程（即将数据输入到一层，再将该层的输出作为输入，输入到下一层），全连接层的输出先经过Relu激活函数再传入下一层。顺便说一句，BP神经网络的隐藏层最多两层，再增加层数并不会提高预测精度。

输入层为28 * 28，因为BP神经网络只支持一维数据，所以要将28 * 28大小的图片转换为28 * 28=784大小的一维向量。

隐藏层维数通常为2的幂次方，这里设为256。

输出层的维数为10，因为这是一个10分类问题，需要最后输出10个概率来判断输入的图片属于哪一类。

3.

device = torch.device("cuda:0")

创建指定设备类型的对象，这里指定设备为GPU。

4.

train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)

加载MNIST数据集，root=" "指定数据集所在的根目录；train=True指定下载的是训练集；transform=transforms.ToTensor()，将图片转换为tensor类型；download=True，如果没有下载数据集就下载数据集。tensor是张量，形式类似于array。为什么要将array转换为tensor？因为要在GPU上加速，数据形式就必须转换为tensor。

test_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())

train=False，表明下载的是测试集。这一句里没有指定download=True，因为pytorch内置了测试集，为了节省资源，不再下载测试集。

5.

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)

定义数据加载器，batch_size指定每一个batch的大小，shuffle=True将数据打乱，随机化。

此时，数据被分成了一个batch一个batch的形式，训练的时候也是以batch为单位送入模型。

6.

model = Network().to(device)

将模型转换为GPU形式

7.

criterion = nn.CrossEntropyLoss()

损失函数为交叉熵损失，通过使用该损失函数，可以顺便在输出层后面自动加softmax，使得最终的输出为每一类的概率。因此在构建模型的时候并没有显示地加入softmax层。

8.

optimizer = optim.Adam(model.parameters())

优化器使用Adam，可以自动调整学习率。什么是优化器？其实就是更新参数的方法，常用的优化器还有随机梯度下降法SGD。

9.

for epoch in range(20):
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 20, i+1, len(train_loader), loss.item()))

这是模型训练的代码。最外层的循环中range(20)表示总共训练20轮。

内层的循环中：

1)

for i, (images, labels) in enumerate(train_loader):

enumerate()是用来读取迭代器的，train_loader是DataLodaer的返回值，是一个以(输入数据,标签)为内容的迭代器。

i 是序号。(images,labels)是包含了图片和标签的元组。注意，这里的(images,labels)是一个batch，而不是单独的图片和标签。

2)

images, labels = images.to(device), labels.to(device)

将数据转换为GPU类型的，这样才能在GPU上运行。

3）

optimizer.zero_grad()

每一轮先将梯度清零，防止使用前几轮累计的梯度更新参数。

4)

outputs = model(images)

将图片输入到模型中，得到输出。

5)

loss = criterion(outputs, labels)

利用输出和label计算loss

6)

loss.backward()

反向传播。

7)

optimizer.step()

使用优化器进行参数更新。

8)

if (i+1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, 20, i+1, len(train_loader), loss.item()))

每经过100个batch，打印已完成的epoch/总的epoch，已完成的batch数/总的batch数，经过的这100个batch中，最后一个batch的loss。

10.

model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy: {:.2f}%'.format(100 * correct / total))

这是预测的代码。

1）

使用model.eval()的原因是：在模型中，我们通常会加上Dropout层和batch normalization层，在模型预测阶段，我们需要将这些层设置到预测模式，model.eval()就是帮我们一键搞定的，如果在预测的时候忘记使用model.eval()，会导致不一致的预测结果。

2)

with语句相当于try-finally。torch.no_grad()用来设置在预测时不对梯度进行更新。

3)

correct = 0
total = 0

correct是预测正确的数量，total是总的数量。

4)

_, predicted = torch.max(outputs.data, 1)

torch.max()返回最大值和最大值的序号，这里不需要最大值，所以写_。使用torch.max()是要得出输出概率中的最大值对应的类别，1表示输出每行的最大值。

5）

total += labels.size(0)

计算总的图片数量。

6）

correct += (predicted == labels).sum().item()

计算预测正确的数量

7）

print('Test Accuracy: {:.2f}%'.format(100 * correct / total))

打印正确率