MegEngine 10 分钟快速上手¶

在 MegStudio 运行	查看源文件

本教程假定你具备一定的 Python 编程基础，并了解深度学习的基础概念。

我们将向你介绍使用 MegEngine 实现的完整的机器学习工作流程，以便你快速地熟悉 MegEngine 常见 Python API 的使用方式。

请先运行下面的代码，检验你的环境中是否已经安装好 MegEngine（安装教程）：

[1]:

import megengine

print(megengine.__version__)

1.4.0

接下来我们将通过 MNIST 手写数字识别的案例帮助你快速上手 MegEngine 的使用。

数据的加载和预处理¶

数据的加载和预处理往往会耗费大量的精力， MegEngine 提供了一系列接口来规范化这些处理工作：

megengine.data.dataset.Dataset 是 MegEngine 中表示数据集的抽象类，存储样本和相应的标签；
megengine.data.DataLoader 负责根据传入的 Dataset 等参数生成一个可迭代的对象。

在 megengine.data.dataset 模块中为用户提供了非常多经典的数据集，比如本次教程中用到的 MNIST 数据集。

注：这一步如果是在国内网络下进行，可能会出现有的压缩包下载失败的情况，这种情况下有两种解决方案：

1.反复运行这段代码几次。

2.手动下载MNIST数据集的几个压缩包到对应文件夹。

[2]:

from megengine.data.dataset import MNIST
import os

# 如果使用 MegStudio 环境，请将 MNIST_DATA_PATH 为 /home/megstudio/dataset/MNIST/
MNIST_DATA_PATH = "/data/datasets/MNIST/"

if not os.path.isdir(MNIST_DATA_PATH):
    os.makedirs(MNIST_DATA_PATH)

# 获取训练数据集，如果本地没有数据集，必须将 download 参数设置为 True或者先手动下载数据集
train_dataset = MNIST(root=MNIST_DATA_PATH, train=True, download=True)
test_dataset = MNIST(root=MNIST_DATA_PATH, train=False, download=True)

15 10:44:42 process the raw files of train set...
100%|██████████████████████████████████| 60000/60000 [00:02<00:00, 29190.52it/s]
100%|████████████████████████████████| 60000/60000 [00:00<00:00, 1284881.83it/s]
15 10:44:45 process the raw files of test set...
100%|██████████████████████████████████| 10000/10000 [00:00<00:00, 29087.64it/s]
100%|████████████████████████████████| 10000/10000 [00:00<00:00, 1359932.56it/s]

对于如何加载自定义的 Dataset ，请参考用户指南使用 Data 处理 I/O 与数据集。

将 Dataset 作为参数传给 DataLoader 时，我们还需要为其指定数据预处理和抽样逻辑：

megengine.data.transfrom 提供了常见的数据变换操作，作为预处理手段，支持 Compose 组合；
megengine.data.sampler 提供了常见的采样方法，如顺序采样和随机采样等，可指定 batch_size 参数。

[3]:

from megengine.data import DataLoader
from megengine.data.transform import ToMode, Pad, Normalize, Compose
from megengine.data.sampler import RandomSampler, SequentialSampler

batch_size=64

# 创建 Sampler
train_sampler = RandomSampler(train_dataset, batch_size=batch_size)
test_sampler = SequentialSampler(test_dataset, batch_size=batch_size)

# 数据预处理方式
transform = Compose([
    Normalize(mean=0.1307*255, std=0.3081*255),
    Pad(2),
    ToMode('CHW'),
])

# 创建 Dataloader
train_dataloader = DataLoader(train_dataset, train_sampler, transform)
test_dataloader  = DataLoader(test_dataset, test_sampler, transform)

[4]:

for X, y in train_dataloader:
    print("Shape of X: ", X.shape, " Data type of X: ", X.dtype) # [N, C, H, W]
    print("Shape of y: ", y.shape, " Data type of y: ", y.dtype)
    break

Shape of X:  (64, 1, 32, 32)
Shape of y:  (64,) int32

定义网络结构¶

在 MegEngine 中定义网络最常见的方式是创建一个继承自 megengine.module.Module 的类，接着：

在 __init__ 中定义网络的层，各类算子可以在 functional 和 module 模块中找到；
通过 forward 方法描述数据通过网络前向传播时依次执行的算子，从而就定义了网络的结构。

[5]:

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

# 定义网络，以经典的 LeNet 结构为例（但激活函数选用 ReLU）
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)

        # 没有参数，实例化这些类以方便在forward函数中调用
        self.relu = M.ReLU()
        self.pool = M.MaxPool2d(2, 2)


    def forward(self, x):
        # 单通道图片, 两层  5x5 卷积 + ReLU + 池化
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        # F.flatten 将原本形状为 (N, C, H, W) 的张量 x 从第一个维度（即 C）开始拉平成一个维度，
        # 得到的新张量形状为 (N, C*H*W) 。 等价于 reshape 操作： x = x.reshape(x.shape[0], -1)
        x = F.flatten(x, 1)
        # 两层全连接 + ReLU
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        # 输出层（分类器）
        x = self.classifier(x)
        return x

# 实例化网络
net = LeNet()
print(net)

LeNet(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5))
  (relu1): ReLU()
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0)
  (conv2): Conv2d(6, 16, kernel_size=(5, 5))
  (relu2): ReLU()
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0)
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (relu3): ReLU()
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (relu4): ReLU()
  (classifier): Linear(in_features=84, out_features=10, bias=True)
)

定义损失函数、优化器¶

为了实现对模型的训练（即对模型中参数的优化），我们还需要定义：

损失函数（Loss Function），大部分常见的损失函数实现在 megengine.function.loss 模块中；
优化器（Optimizer），常见的优化器实现在 megengine.optimizer, 且支持不同的优化策略；
MegEngine 的自动求导功能由 megengine.autodiff 模块实现，其中 GradManager 负责管理梯度。

[6]:

from megengine.optimizer import SGD
from megengine.autodiff import GradManager

optimizer = SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

gm = GradManager().attach(net.parameters())

模型训练¶

MegEngine 中的计算默认以张量（Tensor）作为基础数据结构，因此记得将我们输入的数据转化为 Tensor.

当检测到 GPU 环境时，MegEngine 会自动将相应的计算过程在 GPU 中执行 （无需手动指定设备） 从而实现加速。
如果需要查询 Tensor 所在设备，可以使用 Tensor.device ;
如果需要改变 Tensor 所在设备，可以使用 Tensor.to 或 functional.copy .

我们设置 10 个训练周期，整个过程中将对训练数据集分批次进行预测，根据反向传播算法更新模型的参数。

在没有GPU加速时，尤其是用个人计算机运行时，一个epoch的运行速度比较慢，请耐心等待。建议安装CUDA(>=10.1), cuDNN(>=7.6)进行加速。

以Kaggle NoteBook为例进行测试，CPU型号为Intel(R) Xeon(R) CPU @ 2.00GHz，GPU型号为Tesla P100 PCIe 16GB.

CPU训练时间：121.9s/epoch

GPU训练时间：67.7s/epoch

[7]:

import numpy as np
import megengine as mge

net.train()
total_epochs = 10
for epoch in range(total_epochs):
    total_loss = 0
    for step, (batch_data, batch_label) in enumerate(train_dataloader):
        batch_data = mge.tensor(batch_data)
        batch_label = mge.tensor(batch_label).astype(np.int32)

        with gm:
            pred = net(batch_data)
            loss = F.loss.cross_entropy(pred, batch_label)
            gm.backward(loss)
            optimizer.step().clear_grad()

        total_loss += loss.numpy().item()
    print("epoch: {}, loss {}".format(epoch, total_loss/len(train_dataset)))

epoch: 0, loss 0.002867442769991855
epoch: 1, loss 0.0009182994486764074
epoch: 2, loss 0.0006611305125678578
epoch: 3, loss 0.0005199052049467961
epoch: 4, loss 0.00046812092686692873
epoch: 5, loss 0.000387949584890157
epoch: 6, loss 0.0003237317308783531
epoch: 7, loss 0.00028708203192800286
epoch: 8, loss 0.00025323729968319337
epoch: 9, loss 0.0002447911872838934

模型的保存与加载¶

在 MegEngine 中通过使用 megengine.save 和 megengine.load 进行模型的保存与加载。

我们首先将训练好的模型的保存到本地：

[8]:

mge.save(net.state_dict(), 'mnist_net.mge')

接着我们可以加载本地的模型文件，在测试集上进行预测，以检测模型的性能。

注：如果需要调用C++进行模型部署和推理，不能仅仅直接调用megengine.save保存，而是调用megengine.jit.trace.dump序列化整个计算图。

[9]:

net = LeNet()
state_dict = mge.load('mnist_net.mge')
net.load_state_dict(state_dict)

net.eval()
correct = 0
total = 0
for idx, (batch_data, batch_label) in enumerate(test_dataloader):
    batch_data = mge.tensor(batch_data)
    batch_label = mge.tensor(batch_label).astype(np.int32)

    pred = net(batch_data)
    loss = F.loss.cross_entropy(pred, batch_label)
    predicted = pred.numpy().argmax(axis=1)
    correct += (predicted == batch_label.numpy()).sum().item()
    total += batch_label.shape[0]
print("correct: {}, total: {}, accuracy: {}".format(correct, total, float(correct) / total))

correct: 9900, total: 10000, accuracy: 0.99

新手入门

深度学习零基础实战教程