PyTorch 入门

MNIST 手写数字识别

1. 导入相关库

import torch
# pytorch中最重要的模块，封装了神经网络相关的函数
import torch.nn as nn
# 提供了一些常用的函数，如softmax
import torch.nn.functional as F
# 优化模块，封装了求解模型的一些优化器，如Adam SGD
import torch.optim as optim
# pytorch 视觉库中提供了一些数据变换的接口
from torchvision import transforms
# pytorch 视觉库提供了加载数据集的接口
from torchvision import datasets
# 用于下载并导入数据集
import torchvision
from torch.utils.data import DataLoader

2. 设置超参数

# 由于使用批量训练的方法，需要定义每批的训练的样本数目
batch_size = 64
# 总共训练迭代的次数
epoch = 1
# 设定初始的学习率
LR = 0.001
# 让torch判断是否使用GPU（若使用GPU环境则会更快）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

3. 获取并装载训练集、测试集

# root 用于指定数据集下载后的存放路径
# train 用于指定数据集下载后载入到训练集or测试集（True:载入训练集 False:载入测试集）
# transform 用于指定导入数据集需要对数据进行的操作（转换为tensor数据类型）
# download 设为True表示数据集由程序自动下载
# 下载训练集
train_dataset = datasets.MNIST(root='./data/',
                               train=True,
                               transform=transforms.ToTensor(),
                               download=True)
# 下载测试集
test_dataset = datasets.MNIST(root='./data/',
                              train=False,
                              transform=transforms.ToTensor(),
                              download=True)

# 使用Dataloader数据迭代器加载数据
# dataset 用于指定我们载入的数据集名称
# shuffle 在装载的过程会将数据随机打乱顺序并进行打包，便于后续训练
# 装载训练集
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
# 装载测试集
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=True)

4. 搭建神经网络

# 卷积神经网络LeNet
# 卷积层使用 torch.nn.Conv2d
# 激活层使用 torch.nn.ReLU
# 池化层使用 torch.nn.MaxPool2d
# 全连接层使用 torch.nn.Linear
class LeNet(nn.Module):
    def __init__(self):
        super().__init__()
        # sequential() 按照顺序排列，要保证相邻层的输入输出大小相匹配

        # 卷积 conv2d(in_channels, out_channels, kernel_size, stride, padding)
        # in_channels 网络输入的通道数
        # out_channels 网络输出的通道数
        # kernel_size 卷积核大小q*q
        # stride 卷积移动步长
        # padding 全零填充

        # 激活函数 relu()

        # 最大池化 maxpool2d(kernel_size, stride)
        # kenel_size 最大池化的窗口大小
        # stride 最大池化窗口的移动步长

        # 全连接层 linear(in_features, out_features)
        # in_features 输入维数
        # out_features 输出维数

        # BN层 batchnorm1d(num_features)
        # num_features 特征维度（等于前一层输出的维度）

        self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2), nn.ReLU(),
                                   nn.MaxPool2d(2, 2))

        self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5), nn.ReLU(),
                                   nn.MaxPool2d(2, 2))

        self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120),
                                 nn.BatchNorm1d(120), nn.ReLU())

        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.BatchNorm1d(84),
            nn.ReLU(),
            # 最后的结果一定要变为 10，因为数字的选项是 0 ~ 9
            nn.Linear(84, 10))

    # 前向传播
    def forward(self, x):
        # 卷积+池化
        x = self.conv1(x)
        # 卷积+池化
        x = self.conv2(x)
        # 对参数实现扁平化，便于后面全连接层输入
        x = x.view(x.size()[0], -1)
        # 全连接
        x = self.fc1(x)
        # 全连接
        x = self.fc2(x)
        return x

# 将网络操作移动到gpu或者cpu
net = LeNet().to(device)
# 定义交叉熵损失函数
criterion = nn.CrossEntropyLoss()
# 定义模型优化器（Adam自适应优化算法）
# 输入模型参数和初始学习率
optimizer = optim.Adam(
    net.parameters(),
    lr=LR,
)

5. 训练模型

def train(epoch):
    sum_loss = 0.0
    # 从迭代器抽取图片和标签
    for i, data in enumerate(train_loader):
        inputs, labels = data
        # 将优化器内部参数梯度归零
        optimizer.zero_grad()
        # 将数据传入网络进行前向运算
        outputs = net(inputs)
        # 计算损失函数值
        loss = criterion(outputs, labels)
        # 反向传播
        loss.backward()
        # 通过梯度做一步参数更新
        optimizer.step()
        # 计算每批次的损失值
        sum_loss += loss.item()
        if i % 100 == 99:
            print('[%d,%d] loss:%.03f' %
                  (epoch + 1, i + 1, sum_loss / 100))
            sum_loss = 0.0

6. 测试模型

def test():
    # 将模型变换为测试模式
    net.eval()
    correct = 0
    for data_test in test_loader:
        images, labels = data_test
        # 测试输入数据得到输出
        output_test = net(images)
        # 找到概率最大的分类
        # torch.max()返回value,index（_表示value,predicted表示index）
        # dim=1表示输出所在行（样本）的最大值  dim=0表示输出所在列（类别）的最大值
        _, predicted = torch.max(output_test, 1)
        # 计算正确分类的数目
        correct += (predicted == labels).sum()
    # 输出测试集中正确分类的数目
    print("correct1: ", correct)
    # 输出测试集分类正确率
    print("Test acc: {0}".format(correct.item() /
                                 len(test_dataset)))
                                 
# 程序入口
if __name__ == '__main__':
    for epoch in range(epoch):
        train(epoch)
        test()