MegEngine 快速上手

注解

  • 本教程假设读者具有最基础的 Python 代码编程经验,以及了解 “深度学习” 基本概念;

  • 在进行下面的步骤之前,请确认你已经按照 如何安装 MegEngine 页面的指示安装好 MegEngine.

    >>> import megengine
    >>> print(megengine.__version__)
    

参见

概览

这份快速上手教程将引导你:1. 使用 MegEngine 框架开发出经典的 LeNet 神经网络模型; 2. 使用它在 MNIST 手写数字数据集上完成训练和测试;3. 将模型用于实际的手写数字分类任务。

在最开始,我们会对 MNIST 数据集有一个最基本的了解,并尝试使用 MegEngine 中的 data 模块完成 MNIST 数据集的获取和预处理,将其划分成训练数据集和测试数据集; 同时你还会准备好相应的 DataLoader, 负责在后续训练和测试模型时完成数据供给。 紧接着你会用 functionalmodule 模块设计好 LeNet 模型结构。 接下来的步骤也很简单,对模型进行训练!训练的过程中我们需要用到 autodiff 模块和 optimizer 模块,前者在训练的过程中会记录梯度信息, 后者负责根据梯度信息对模型中的参数进行更新,以达到优化的目的。 最终,我们会对训练好的模型进行测试,你也可以用自己的手写数字样本来试试效果~

注意:本教程的目的是为 MegEngine 初见用户展示框架最基本的使用流程, 因此不会对每个步骤以及背后的原理进行非常详细的解释,也不会展现出 MegEngine 的全部特性。 如果你对整个深度学习的流程不是很清楚,不用担心,可以尝试跑通这个教程,最后会有进一步的指引。

数据加载和预处理

数据集介绍

MNIST 1 手写数字数据集中包含 60,000 张训练图像和 10,000 张测试图像,每张图片是 28x28 像素的灰度图。 如今 MNIST 已然成为机器学习领域的 “Hello, world!”, 用来验证框架和库的可用性。

../_images/MnistExamples.png

By Josef Steppan - Own work , CC BY-SA 4.0

获取 MNIST 数据集

在 MegEngine 中可以 使用已经实现的数据集接口 来获取 MNIST 数据集:

from megengine.data.dataset import MNIST

MNIST_DATA_PATH = "/data/datasets/MNIST"

train_dataset = MNIST(MNIST_DATA_PATH, train=True)
test_dataset = MNIST(MNIST_DATA_PATH, train=False)

准备 DataLoader

我们将上一步得到的训练集和测试集作为参数输入给 DataLoader:

import megengine.data as data
import megengine.data.transform as T

train_sampler = data.RandomSampler(train_dataset, batch_size=64)
test_sampler = data.SequentialSampler(test_dataset, batch_size=4)

transform = T.Compose([
    T.Normalize(0.1307*255, 0.3081*255),
    T.Pad(2),
    T.ToMode("CHW"),
])

train_dataloader = data.DataLoader(train_dataset, train_sampler, transform)
test_dataloader = data.DataLoader(test_dataset, test_sampler, transform)

在上面的代码中,我们对数据集的抽样规则和预处理策略也进行了定义, 例如指定了训练集的 batch_size 为 64, 抽样方式为随机抽样… 并分别将对应的 samplertransform 作为构造 DataLoader 的初始化参数提供。

参见

想要了解更多细节,可以参考 使用 Data 构建输入 Pipeline

定义模型结构

LeNet 2 网络模型的结构如下图所示(图片截取自论文):

../_images/lenet5.png

Architecture of LeNet a Convolutional Neural Network here for digits recognition. Each plane is a feature map ie a set of units whose weights are constrained to be identical.

在 MegEngine 中定义网络最常见的方式是创建一个继承自 Module 的类:

import megengine.functional as F
import megengine.module as M

class LeNet(M.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = M.Conv2d(1, 6, 5)
        self.conv2 = M.Conv2d(6, 16, 5)
        self.fc1 = M.Linear(16 * 5 * 5, 120)
        self.fc2 = M.Linear(120, 84)
        self.classifier = M.Linear(84, 10)

        self.relu = M.ReLU()
        self.pool = M.MaxPool2d(2, 2)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = F.flatten(x, 1)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.classifier(x)
        return x


model = LeNet()
  • 需要在 __init__ 方法中调用 super().__init__;

  • 需要在 __init__ 方法中定义需要用到的结构,并在 forward 中定义前向计算过程。

参见

想要了解更多细节,可以参考 使用 Module 定义模型结构

训练:优化模型参数

得到前向计算输出后,为了优化模型参数,我们还需要:

  • 使用 GradManager 对参数梯度进行管理;

  • 使用 Optimizer 进行反向传播和参数更新(以 SGD 为例)。

import megengine.optimizer as optim
import megengine.autodiff as autodiff

gm = autodiff.GradManager().attach(model.parameters())
optimizer = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    weight_decay=5e-4
)

接下来训练我们的模型:将训练数据集分批地喂入模型,前向计算得到预测值, 根据设计好的损失函数(本教程中使用交叉熵 cross_entropy )计算。 接着调用 GradManager.backward 方法来自动进行反向计算并记录梯度信息, 然后根据这些梯度信息来更新模型中的参数,即调用 Optimizer.step 方法。

epochs = 10
model.train()
for epoch in range(epochs):
    total_loss = 0
    for batch_data, batch_label in train_dataloader:
        batch_data = megengine.Tensor(batch_data)
        batch_label = megengine.Tensor(batch_label)

        with gm:
            logits = model(batch_data)
            loss = F.nn.cross_entropy(logits, batch_label)
            gm.backward(loss)
            optimizer.step().clear_grad()

        total_loss += loss.item()

    print(f"Epoch: {epoch}, loss: {total_loss/len(train_dataset)}")

警告

记得将数据转为 MegEngine Tensor 格式,参考 深入理解 Tensor 数据结构

参见

想要了解更多细节,可以参考 Autodiff 基本原理与使用 / 使用 Optimizer 优化参数

测试:评估模型性能

在测试集上验证一下我们刚才训练好的 LeNet 模型的性能:

model.eval()
correct, total = 0, 0
for batch_data, batch_label in test_dataloader:
    batch_data = megengine.Tensor(batch_data)
    batch_label = megengine.Tensor(batch_label)

    logits = model(batch_data)
    pred = F.argmax(logits, axis=1)
    correct += (pred == batch_label).sum().item()
    total += len(pred)

print(f"Correct: {correct}, total: {total}, accuracy: {float(correct)/total}")

通常会得到一个在测试集上接近甚至超过 99% 预测正确率的模型。

注:通常的训练流程中应当使用验证集,每训练一段时间就及时验证,这里简化了这一步。

推理:用单张图片验证

我们也可以选择使用自己的手写数字图片来验证模型效果(你可以选择使用自己的图片):

../_images/handwrittern-digit.png
import cv2
import numpy as np

def process(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    image = cv2.resize(image, (32, 32))
    image = np.array(255 - image)
    return image

image = cv2.imread("/data/handwrittern-digit.png")
processed_image = process(image)

上面是针对输入图片样本所做的一些必要预处理步骤,接下来将其输入模型进行推理:

>>> logit = model(megengine.Tensor(processed_image).reshape(1, 1, 32, 32))
>>> pred = F.argmax(logit, axis=1).item()
>>> pred
6

可以发现,我们训练出的 LeNet 模型成功地将手写该数字图片的标签类别预测为 6 !

参见

这里展示的是最简单的模型推理情景,MegEngine 是一个训练推理一体化的框架, 能将训练好的模型导出,在 C++ 环境下高效地进行推理部署,可参考 模型部署总览与流程建议 中的介绍。

接下来做些什么?

我们已经成功地使用 MegEngine 框架完成了手写数字分类任务,很简单吧~

文档中还提供了更多内容

如果你对整个机器学习(深度学习)的流程依旧不是很清楚,导致阅读本教程有些吃力,不用担心。 我们准备了更加基础的 《 MegEngine 深度学习入门教程 》—— 它可以看作是当前教程内容的手把手教学版本,补充了更多细节和概念解释。 将从机器学习的基本概念开始讲起,循序渐进地帮助你理解整个开发流程, 在接触到更多经典模型结构的同时,也会更加了解如何使用 MegEngine 框架。 一些像是 保存与加载模型(S&L)使用 Hub 发布和加载预训练模型 的用法,也会在该系列教程中进行简单介绍。

同时,由于这仅仅是一份快速上手教程,许多模型开发的进阶特性没有进行介绍,例如 分布式训练(Distributed Training) / 量化(Quantization) … 等专题,可以在 用户指南 中找到。 值得一提的是,MegEngine 不仅仅是一个深度学习训练框架,同时也支持便捷高效的模型推理部署。 关于模型推理部署的内容,可以参考 模型部署总览与流程建议 页面的介绍与 《 MegEngine 模型推理部署教程 》。

任何人都可以成为 MegEngine 教程的贡献者

由于开发者视角所带来的一些局限性,我们无法做到完全以用户视角来撰写文档中的各块内容,尽善尽美是长期追求。 如果你在阅读 MegEngine 教程的过程中产生了疑惑,或是有任何的建议,欢迎你加入 MegEngine 文档建设中来。

参考 如何为文档做贡献 页面了解更多细节。

参考文献

1

Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2010.

2

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.