MegEngine 10 分钟快速上手

image0 在 MegStudio 运行

image1 查看源文件

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

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

请先运行下面的代码,检验你的环境中是否已经安装好 MegEngine(安装教程):

[1]:
import megengine

print(megengine.__version__)
1.4.0

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

数据的加载和预处理

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

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

[2]:
from megengine.data.dataset import MNIST

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

# 获取训练数据集,如果本地没有数据集,请将 download 参数设置为 True
train_dataset = MNIST(root=MNIST_DATA_PATH, train=True, download=False)
test_dataset = MNIST(root=MNIST_DATA_PATH, train=False, download=False)
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 ,请参考用户指南 利用 Dataset 封装一个数据集

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) # [N, C, H, W]
    print("Shape of y: ", y.shape, y.dtype)
    break
Shape of X:  (64, 1, 32, 32)
Shape of y:  (64,) int32

定义网络结构

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

  • __init__ 中定义网络的层,各类算子可以在 functionalmodule 模块中找到;

  • 通过 forward 方法描述数据通过网络前向传播时依次执行的算子,从而就定义了网络的结构。

[5]:
import megengine.module as M
import megengine.functional as F

# 定义网络,以经典的 LeNet 结构为例(但激活函数选用 ReLU)
class LeNet(M.Module):
    def __init__(self):
        super().__init__()
        # 单信道图片, 两层  5x5 卷积 + ReLU + 池化
        self.conv1 = M.Conv2d(1, 6, 5)
        self.relu1 = M.ReLU()
        self.pool1 = M.MaxPool2d(2, 2)
        self.conv2 = M.Conv2d(6, 16, 5)
        self.relu2 = M.ReLU()
        self.pool2 = M.MaxPool2d(2, 2)
        # 两层全连接 + ReLU
        self.fc1 = M.Linear(16 * 5 * 5, 120)
        self.relu3 = M.ReLU()
        self.fc2 = M.Linear(120, 84)
        self.relu4 = M.ReLU()
        # 分类器
        self.classifier = M.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(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)
        x = self.relu3(self.fc1(x))
        x = self.relu4(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)
)

定义损失函数、优化器

为了实现对模型的训练(即对模型中参数的优化),我们还需要定义:

[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.tofunctional.copy .

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

[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.savemegengine.load 进行模型的保存与加载。

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

[8]:
mge.save(net.state_dict(), 'mnist_net.mge')

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

[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