PyTorch官方例程详解（1.使用卷积进行图像分类）

本文属于PyTorch官方教程解析计划系列文章，旨在对官方教程文档和例程代码进行详细解析，加快新手入门过程

 PyTorch的官方Tutoral页面提供了14个例程供大家学习，本次对第一个例程进行详解

例程地址与该例程描述如下：

PyTorch Examples — PyTorchExamples 1.11 documentatin

 代码地址如下：examples/mnist at main · pytorch/examples (github.com)

项目结构如图所示，可以看到非常简洁，仅运行main.py即可：

完整版代码如下：

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))


def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--no-mps', action='store_true', default=False,
                        help='disables macOS GPU training')
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    use_mps = not args.no_mps and torch.backends.mps.is_available()

    torch.manual_seed(args.seed)

    if use_cuda:
        device = torch.device("cuda")
    elif use_mps:
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': 1,
                       'pin_memory': True,
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

    if args.save_model:
        torch.save(model.state_dict(), "mnist_cnn.pt")


if __name__ == '__main__':
    main()

分析：

可以看到文件由五部分组成，分别是：导库部分、网络部分、训练部分、测试部分、主函数部分，为保证代码规范，我们在使用相关函数的时候可直接模仿该形式，现对其分别进行分析如下：

1.导库部分

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR

from __future__ import print_function：

                保证代码对Python2的兼容性，可以直接放在文件开头，自己复现代码时可以不加
import argparse：

                argparse 是 Python 内置的一个用于命令项选项与参数解析的模块
import torch：

                使用PyTorch框架
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim：

                优化函数
from torchvision import datasets, transforms

                PyTorch自带数据集和图像处理函数
from torch.optim.lr_scheduler import StepLR：

                PyTorch提供了几种调整学习率的方法，如下表所示，通过在迭代过程中调整学习率，将学习率由大到小进行调整，可以达到的效果，在本例程中使用StepLR进行调整：

lr_scheduler.LambdaLR	将初始学习率乘以人工规则所生成的系数λ
lr_scheduler.MultiplicativeLR	将每个参数组的学习率乘以指定函数中给出的因子。
lr_scheduler.StepLR	每 step_size 个时期通过 gamma 衰减每个参数组的学习率。
lr_scheduler.MultiStepLR	一旦 epoch 的数量达到一定数值，就通过 gamma 衰减每个参数组的学习率。
lr_scheduler.ConstantLR	将每个参数组的学习率衰减一个小的常数因子，直到 epoch 的数量达到预定义的里程碑：total_iters。
lr_scheduler.LinearLR	通过线性改变小的乘法因子来衰减每个参数组的学习率，直到 epoch 的数量达到预定义的里程碑：total_iters。
lr_scheduler.ExponentialLR	每个 epoch 通过 gamma 衰减每个参数组的学习率。
lr_scheduler.PolynomialLR	在给定的 total_iters 中使用多项式函数衰减每个参数组的学习率。
lr_scheduler.CosineAnnealingLR	Set the learning rate of each parameter group using a cosine annealing schedule, where \eta_{max}ηmax​ is set to the initial lr and T_{cur}Tcur​ is the number of epochs since the last restart in SGDR:
lr_scheduler.ChainedScheduler	学习率调度器的链表。
lr_scheduler.SequentialLR	接收预计在优化过程中按顺序调用的调度程序列表和里程碑点，这些里程碑点提供准确的时间间隔以反映应该在给定时期调用哪个调度程序。
lr_scheduler.ReduceLROnPlateau	当指标停止改善时降低学习率。
lr_scheduler.CyclicLR	根据循环学习率策略 (CLR) 设置每个参数组的学习率。
lr_scheduler.OneCycleLR	按照1cycle学习率策略设置各参数组的学习率。
lr_scheduler.CosineAnnealingWarmRestarts使用余弦退火计划设置每个参数组的学习率

 2.模型定义部分

 PyTorch提供了用来定义模型的类nn.Module，在使用PyTorch自定义模型的时候，需要继承nn.Module类并重写__init__和forward函数，其中，__init__函数用来进行初始化操作，包括使用外部参数初始化、模型各层的定义，forward函数用来进行前向传播。

训练部分：

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
            if args.dry_run:
                break

① 将模型设置为训练模式，启用模型中的 Batch Normalization 和 Dropout，如无可省略

② 每次从train_loader中读取一批数据

③ 将数据移入设备（多用于在电脑中有GPU时将数据移入显存，如使用CPU计算，则可省略）

④ 梯度清零

⑤ 使用损失函数计算损失

⑥ 更新所有的参数

⑦打印训练过程百分比和Loss

测试部分：

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

①：将模型设置为评估模式

②：建立两个变量用来保存测试的loss和准确率

③每次从test_loader中取一组数据和标签

        分别存到指定设备，如有GPU，则保存到显存中，如无，则可省略

        使用模型对数据进行预测

        将预测值与标签放入损失函数，计算本次损失，并加入总的测试损失中

        预测结果是模型输出里面最大的

计算平均损失，并输出有关信息

主函数与其他部分，此处我们进行拆分讲解

主函数内部①：参数设置部分，该部分代码作用是通过命令行来设置模型参数

# main函数内部①：

    
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')

    # batch-size：设置批量大小
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    # test-batch-size：设置测试时的批量大小
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')

    # epoch:：设置迭代次数
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')

    # lr:设置学习率
    parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
                        help='learning rate (default: 1.0)')

    # gamma：设置参数gamma
    parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
                        help='Learning rate step gamma (default: 0.7)')

    # no-cuda：设置不使用英伟达显卡进行加速
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')

    # no-mps：设置不使用基于M系列内核的苹果电脑的GPU进行加速
    parser.add_argument('--no-mps', action='store_true', default=False,
                        help='disables macOS GPU training')

    # dry-run：设置仅运行一次来测试程序是否能正确运行
    parser.add_argument('--dry-run', action='store_true', default=False,
                        help='quickly check a single pass')

    # seed：设置随机数种子
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')

    # log-interval
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    # save-model：设置保存模型
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    args = parser.parse_args()
 

argparse的出现形式：

1：parser = argparse.ArgumentParser()  # 创建解析对象

2：parser.add_argument() # 添加解析参数

3：parser.parse_args() # 调用parse_args()方法进行解析

命令行使用方法：

在命令行中输入以下代码，可以看到代码中定义的帮助信息如下图所示

python3 main.py --batch-size 64 --test-batch-size 64 --epochs 2 --lr 0.01 --gamma 0.5 --dry-run --seed 10 --log-interval 1 --save-model

 运行时通过在命令行中按照提示信息输入对应参数即可，示例如下：

python3 main.py --batch-size 64

即可将batch_size设为64，其他参数同理

继续分析main函数其他部分

# main函数内部②

    use_cuda = not args.no_cuda and torch.cuda.is_available()
    use_mps = not args.no_mps and torch.backends.mps.is_available()

    torch.manual_seed(args.seed)

    if use_cuda:
        device = torch.device("cuda")
    elif use_mps:
        device = torch.device("mps")
    else:
        device = torch.device("cpu")

    train_kwargs = {'batch_size': args.batch_size}
    test_kwargs = {'batch_size': args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {'num_workers': 1,
                       'pin_memory': True,
                       'shuffle': True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

CUDA的相关设置

①：对于CUDA来说，命令行中设置的开关（默认使用CUDA）具有最高优先级，如果未手动关闭，且系统检测到CUDA可用，则use_cuda有效，表示使用CUDA，Macbook的mps同理 

②：随机数种子设置，一个确定的随机数种子产生的随机数也是确定的，方便模型的复现

③：若use_cuda有效，则PyTorch会在训练及测试的过程中将有关模型及代码转入显存，mps同理，若cuda与mps均不可用，则使用CPU进行计算，使用CPU计算可以节省数据转入显存的时间，但计算时间会超过GPU（在模型较小时差距并不明显，但随着模型参数及数据量的增加，差距会明显加大）

# main函数内部③
 

transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])
    dataset1 = datasets.MNIST('../data', train=True, download=True,
                       transform=transform)
    dataset2 = datasets.MNIST('../data', train=False,
                       transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

transform部分：torchvision.transforms是pytorch中的图像预处理包。一般用Compose把多个步骤整合到一起，此处先使用ToTensor函数将输入的图像数据转为PyTorch使用的tensor格式张量，transforms.Normalize的两个参数为均值和方差，此处手动对均值和方差进行设置

dataset部分：此处使用PyTorch内置的MNIST手写数字数据集，若之前没有下载过，则将download设置为True，程序执行时会开始进行下载，并将数据集放在上级文件夹的data文件夹中，若之前下载过，则可将download设置为False，手动在上级文件夹中创建data文件夹，并将数据集文件夹移入即可。使用内置数据集作为训练集时，需要将train参数设置为True，作为测试集时，将train设置为False即可。

dataloader部分：dataloader是数据集的加载器，其接收PyTorch封装好的dataset作为第一个参数，每次向模型输入一批数据和标签（测试集没有标签）

# main函数内部，4

model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
    for epoch in range(1, args.epochs + 1):
        train(args, model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
        scheduler.step()

    if args.save_model:
        torch.save(model.state_dict(), "mnist_cnn.pt")

①训练开始前，将模型移入对应的设备（GPU等）

②指定优化器：torch.optim中提供了许多优化函数，已列入表中，例程中使用StepLR优化函数，其中的gamma参数通过命令行传入

Adadelta	使用 Adadelta 算法
Adagrad	使用 Adagrad 算法
Adam	使用 Adam 算法
AdamW	使用 AdamW 算法
SparseAdam	使用适用于稀疏张量的 Adam 算法的惰性版本
Adamax	使用Adamax 算法(基于无穷范数的 Adam 变体).
ASGD	使用平均随机梯度下降(Averaged Stochastic Gradient Descent)
LBFGS	使用L-BFGS 算法
NAdam	使用NAdam 算法.
RAdam	使用RAdam 算法.
RMSprop	使用RMSprop 算法.
Rprop	使用弹性反向传播算法.
SGD	使用随机梯度下降算法 (可选momentum).

③scheduler：设置学习率下降算法

④开始遍历：在每次遍历的过程中调用train和test函数,并进行一次学习率变化

⑤保存模型：如果设置save_model参数为True，则在迭代完成后保存模型

代码的最后一步：

if __name__ == '__main__':
    main()

如果该文件本身作为工程运行，则执行main函数，否则，不执行main函数，提供组件供其他文件中的函数调用。

总结：

该例程的特点如下：

1.提供命令行设置参数的接口，方便在命令行中修改参数运行

2.使用print输出的部分封装在train和test中，将main函数简化

3.使用层级式的前向传播，这种方式适合模型层数较少的情况，若层数增加，使用nn.Sequential更加适宜

4.支持使用M系列内核的MacBook的GPU加速

参考资料：

(45条消息) python之parser.add_argument()用法——命令行选项、参数和子命令解析器_夏普通的博客-CSDN博客_parser.add

(45条消息) 使用PyTorch时，最常见的4个错误_人工智能与算法学习的博客-CSDN博客

(45条消息) transforms.Compose()函数_马鹏森的博客-CSDN博客_transforms.compose

最后：

受作者水平限制，文章可能有内容描述不够准确恰当，欢迎在评论中指出，文章将持续更新

谢谢！