1. 前言

Megengine 是旷视科技开发的一款训练推理一体化的深度学习框架，类似于 pytorch ，tensorflow 。

使用 Megengine 可以快速实现常见的深度学习模型，本文将使用 Megengine 实现手写数字识别，以完成深度学习的两大步骤：训练和预测。通过本文，读者对深度学习的最基本流程和 Megengine 框架的使用方法有大致了解。

2. 环境安装

在命令行输入下列语句即可安装 Megengine ，建议使用 python 版本为 3.5 到 3.8

python3 -m pip install --upgrade pip python3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html 

安装完成后可以在命令行测试是否安装成功。

python3 import megengine print(megengine.__version__) 

3. 训练

本部分训练代码来自 Megengine 官方教程，需要详细了解细节请前往 MegEngine 快速上手

3.1 数据集准备

3.1.1 下载数据集

深度学习的第一步为准备数据集，通常会为数据集写一个接口来访问数据集，并对数据进行预处理。

Megengine 中已经实现了 MNIST 数据集的接口，我们可以通过以下代码直接获取。如果想要制作或使用其他数据集，可以点击这里进行学习。

from megengine.data.dataset import MNIST DATA_PATH = "./datasets/MNIST" #第一次运行后，将 download 改为 False train_dataset = MNIST(DATA_PATH, train=True, download=True) test_dataset = MNIST(DATA_PATH, train=False, download=True) 

3.1.2 数据加载及预处理

上面使用 MNIST （）完成数据集的加载和 Dataset 的构建，接下来将对数据进行加载，修改数据。使用 DataLoader 、Sampler 和 Transform 实现。

DataLoader

功能： 构建可迭代的数据装载器，非常灵活地从数据集连续获取小批量数据 参数：

• dataset 需要从中分批加载的数据集。
• sampler (Optional) 定义从数据集中采样数据的策略。
• transform (Optional) 定义抽样批次的转换策略。对数据需要作的变换 默认：None
• ...

RandomSampler

功能：创建一个列表，包含所有数据的索引，可实现数据的随机取样 参数：

• dataset 待采样的目标数据集。
• batch_size 使用 batch 方法时指定 batch 大小。
• drop_last 如果 batch 大小不能整除数据集大小时，为 True 则放弃最后一个不完整的 batch; 为 False 则最后一个 batch 可能比较小。默认：False
• ...

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) 

3.2 模型

接下来定义网络结构，LeNet 的网络结构如下图所示。 [图片上传失败...(image-41aaca-1660811354056)] 定义网络结构主要为两步：定义网络子模块和连接网络子模块。如下代码所示，使用 init 方法创建子模块，forward()方法连接子模块。

import megengine.functional as F import megengine.module as M class LeNet(M.Module): def __init__(self): super().__init__() #输入大小为(batch, 1, 32, 32),输出大小为(batch, 6, 28, 28) 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.fc3 = M.Linear(84, 10) def forward(self, x): x = F.max_pool2d(F.relu(self.conv1(x)), 2) x = F.max_pool2d(F.relu(self.conv2(x)), 2) x = F.flatten(x, 1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x 

• class LeNet 继承自 Module 的类,想要了解更多关于构建模型的细节，可以参考参考 module 定义模型结构

3.3 训练准备

• GradManager 负责计算梯度，管理需要优化的参数，用 attach()方法添加，加入的参数才会被计算保存梯度值。对梯度计算了解更多请点击 Autodiff 基本原理与使用
• optimizer 选择使用的优化方法，这里使用 Optimizer.SGD （随机梯度下降法）作为优化器，对优化器了解更多请点击使用 optimizer 优化参数

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 ) 

3.4 训练迭代

接下来进入程序的主逻辑，开始训练模型。使用两个嵌套循环，一个大循环为一个 epoch ，遍历一次数据集，计算一次准确度。

每个小循环为一个 batch ，将一批数据传入模型中，进行前向计算得到预测概率，使用交叉熵(cross_entropy)来计算 loss, 接着调用 GradManager.backward 方法进行反向计算并且记录每个 tensor 的梯度信息。然后使用 Optimizer.step 方法更新模型中的参数。由于每次更新参数后不自动清除梯度，所以还需要调用 clear_grad 方法。

import megengine 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)}") 

3.5 保存模型

常用的神经网络都具有非常大数量级的参数，每次训练需要花费很长时间，为了能够训练中断后能够按照上次训练的成果接着训练，我们可以每 10 个 epoch 保存一次模型（或更多）。保存模型有几种方法，如表所示。方法详细介绍请点击保存与加载模型。

方法	优劣
保存 /加载整个模型	任何情况都不推荐
保存加载模型状态字典	适用于推理，不满足恢复训练要求
保存加载检查点	适用于推理或恢复训练
导出静态图模型	适用于推理，追求高性能部署
我们选择保存加载检查点，既可以用于恢复训练也可以推理。保存时调用 megengine.save()方法，参数如下：

megengine.save({ "epoch": epoch, "state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "loss": loss, ... }, PATH) 

然后就可以愉快的进行训练了，观察训练结果，当 loss 下降到一定地步，准确率满足要求后，终止训练.

如果训练发生中断，可以调用 load()方法和 optimizer.load_state_dict()方法，对模型的加载，重新开始训练。代码如下：

model = LeNet() optimizer = optim.SGD() checkpoint = megengine.load(PATH) model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) epoch = checkpoint["epoch"] loss = checkpoint["loss"] model.eval() # - or - model.train() 

4. 推理

上面几个章节已经完成深度学习大部分内容，已经能够产生一个需要的算法模型。这个算法对准备好的数据集有比较好的拟合效果，但是我们的最终目的是用模型进行推理，即能够对新的数据进行预测。这将是下面介绍的内容。

首先有一种很简的方法，使用 python 加载模型并设定 model.eval()，代码如下所示，这样就可以简单调用训练好的模型用以实际。

from train import LeNet import cv2 import numpy as np import megengine import megengine.data.transform as T import megengine.functional as F IMAGE_PATH = "./test.png" CHECK_POINT_PATH = "./checkpoint.pkl" def load_model(check_point_path = CHECK_POINT_PATH): model = LeNet() check_point = megengine.load(check_point_path) #注意 checkpoint 保存时模型对应的键，此处为 state_dict model.load_state_dict(check_point["state_dict"]) model.eval() return model def main(): # 加载一张图像为灰度图 image = cv2.imread(IMAGE_PATH,cv2.IMREAD_GRAYSCALE) image = cv2.resize(image, (32, 32)) #将图片变换为黑底白字 image = np.array(255-image) tensor_image = megengine.tensor(image).reshape(1, 1, 32, 32) model = load_model() logit= model(tensor_image) pred = F.argmax(logit, axis=1).item() print("number:", pred) if __name__ == "__main__": main() 

不过在实际部署中，还需要考虑部署环境，推理速度等因素，所以从训练好模型到部署落地还有很长的路。Megengine 由于其设计特点训练推理一体化，可以方便地将训练模型部署。这将是下一章介绍的内容，下一章将使用 C++ 调用 Megengine lite ，进行高效部署。

参考文献

[1]: MegEngine 快速上手 [2]: Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist, 2010. [3]: Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):22782324, 1998.

附录：train.py

from megengine.data.dataset import MNIST from megengine import jit, tensor import megengine import numpy as np import megengine.data as data import megengine.data.transform as T import megengine.functional as F import megengine.module as M import megengine.optimizer as optim import megengine.autodiff as autodiff DATA_PATH = "./datasets/train/" def load_data(data_path =DATA_PATH): train_dataset = MNIST(DATA_PATH) test_dataset = MNIST(DATA_PATH) train_sampler = data.RandomSampler(train_dataset, batch_size=64) test_sampler = data.SequentialSampler(test_dataset, batch_size=2) 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) return train_dataloader, test_dataloader #Define model 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.fc3 = M.Linear(84, 10) def forward(self, x): x = F.max_pool2d(F.relu(self.conv1(x)), 2) x = F.max_pool2d(F.relu(self.conv2(x)), 2) x = F.flatten(x, 1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x def train(dataloader): model = LeNet() # GradManager and Optimizer setting gm = autodiff.GradManager().attach(model.parameters()) optimizer = optim.SGD( model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4 ) # Training and validation nums_epoch = 50 for epoch in range(nums_epoch): training_loss = 0 nums_train_correct, nums_train_example = 0, 0 nums_val_correct, nums_val_example = 0, 0 for step, (image, label) in enumerate(dataloader[0]): image = megengine.Tensor(image) label = megengine.Tensor(label) with gm: score = model(image) loss = F.nn.cross_entropy(score, label) gm.backward(loss) optimizer.step().clear_grad() training_loss += loss.item() * len(image) pred = F.argmax(score, axis=1) nums_train_correct += (pred == label).sum().item() nums_train_example += len(image) training_acc = nums_train_correct / nums_train_example training_loss /= nums_train_example for image, label in dataloader[1]: image = megengine.Tensor(image) label = megengine.Tensor(label) pred = F.argmax(model(image), axis=1) nums_val_correct += (pred == label).sum().item() nums_val_example += len(image) val_acc = nums_val_correct / nums_val_example #每十次 epoch 保存一次模型 if epoch%2 == 0: megengine.save( {"epoch":epoch, "state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "loss": loss, }, "./checkpoint.pkl") print(f"Epoch = {epoch}, " f"train_loss = {training_loss:.3f}, " f"train_acc = {training_acc:.3f}, " f"val_acc = {val_acc:.3f}") def dumpy_mge(pkl_path = "checkpoint.pkl"): model = LeNet() check_point = megengine.load(pkl_path) model.load_state_dict(check_point["state_dict"]) model.eval() @jit.trace(symbolic=True, capture_as_cOnst=True) def infer_func(input, *, model): pred = model(input) return pred input = megengine.Tensor(np.random.randn(1, 1, 32, 32)) output = infer_func(input, model=model) infer_func.dump("./lenet.mge", arg_names=["input"]) if __name__=='__main__': train(load_data()) 

C++ 推理

前半部分我们完成了深度学习的训练，得到了 LeNet 训练权重文件，后面我们将使用训练权重文件导出静态图模型。并使用 C++调用模型完成实际部署的模拟。

准备工作

在上一章中，我们提到有四种保存模型的方法，如下表所示，为了训练方便起见，保存了 checkpoint 文件。但实际部署中我们经常使用静态图模型，所以我们首先要完成静态图导出。

方法	优劣
保存 /加载整个模型	任何情况都不推荐
保存加载模型状态字典	适用于推理，不满足恢复训练要求
保存加载检查点	适用于推理或恢复训练
导出静态图模型	适用于推理，追求高性能部署

到处静态图在 megengine 中有较完整的教程，请参考导出序列化模型文件（ Dump ）。主要分为三步:

1. 将循环内的前向计算、反向传播和参数优化代码提取成单独的函数，如下面例子中的 train_func()
2. 将网络所需输入作为训练函数的参数，并返回任意你需要的结果（如输出结果、损失函数值等）；
3. 用 jit 模块中的 trace 装饰器来装饰这个函数，将其中的代码变为静态图代码。

在上一章最后的附录 train.py 中有 dump 静态图的方法，代码如下：

 from megengine import jit def dump_mge(pkl_path = "checkpoint.pkl"): model = LeNet() check_point = megengine.load(pkl_path) model.load_state_dict(check_point["state_dict"]) model.eval() @jit.trace(symbolic=True, capture_as_cOnst=True) def infer_func(input, *, model): pred = model(input) pred_normalized = F.softmax(pred) return pred_normalized input = megengine.Tensor(np.random.randn(1, 1, 32, 32)) output = infer_func(input, model=model) infer_func.dump("./lenet.mge", arg_names=["input"]) 

调用 dump_mge 方法即可完成静态图导出。

inference 代码

代码的主逻辑为：

1. 创建 Network
2. 使用 load_model()载入模型
3. 使用 stb 预处理图片（加载和 resize)，然后归一化，载入进 input tensor
4. 使用 network->forward()和 network->wait()完成推理逻辑。
5. 获取模型输出 tensor ，并对其进行处理。

推理代码为：

//inference.cpp #include <iostream> #include <stdlib.h> #define STB_IMAGE_IMPLEMENTATION #include "stb/stb_image.h" #define STB_IMAGE_WRITE_IMPLEMENTATION #include "stb/stb_image_write.h" #define STB_IMAGE_RESIZE_IMPLEMENTATION #define STB_IMAGE_RESIZE_STATIC #include "stb/stb_image_resize.h" #include "lite/network.h" #include "lite/tensor.h" //注意在这里修改测试图片与所用模型 #define IMAGE_PATH "./test.png" #define MODEL_PATH "./lenet.mge" void preprocess_image(std::string pic_path, std:shared_ptr<lite::Tensor> tensor) { int width, height, channel; uint8_t* image = stbi_load(pic_path.c_str(), &width, &height, &channel, 0); printf("Input image %s with height=%d, width=%d, channel=%d\n", pic_path.c_str(), width, height, channel); auto layout = tensor->get_layout(); auto pixels = layout.shapes[2] * layout.shapes[3]; size_t image_size = width * height * channel; size_t gray_image_size = width * height * 1; unsigned char *gray_image = (unsigned char *)malloc(gray_image_size); for(unsigned char *p=image, *pg=gray_image; p!=image+image_size; p+=channel,pg++) { *pg = uint8_t(*p + *(p+1) + *(p+2))/3.0; } //! resize to tensor shape std::shared_ptr<std::vector<uint8_t>> resize_int8 = std::make_shared<std::vector<uint8_t>>(pixels * 1); stbir_resize_uint8( gray_image, width, height, 0, resize_int8->data(), layout.shapes[2], layout.shapes[3], 0, 1); free(gray_image); stbi_image_free(image); //! 减去均值，归一化 unsigned int sum = 0; for(unsigned char *p=gray_image; p!=gray_image+gray_image_size;p++){ sum += *p; } sum /= gray_image_size; float* in_data = static_cast<float*>(tensor->get_memory_ptr()); for (size_t i = 0; i < pixels; i++) { in_data[i] = resize_int8->at(i)-sum; } } int main() { //创建网络 std::shared_ptr<lite::Network> network = std::make_shared<lite::Network>(); //加载模型 network->load_model(MODEL_PATH); std::shared_ptr<lite::Tensor> input_tensor = network->get_io_tensor("input"); preprocess_image(IMAGE_PATH, input_tensor); //将图片转为 Tensor network->forward(); network->wait(); std::shared_ptr<lite::Tensor> output_tensor = network->get_output_tensor(0); float* predict_ptr = static_cast<float*>(output_tensor->get_memory_ptr()); float max_prob = predict_ptr[0]; size_t number = 0; //寻找最大的标签 for(size_t i=0; i<10; i++) { float cur_prob = predict_ptr[i]; if(cur_prob>max_prob) { max_prob = cur_prob; number = i; } } std::cout << "the predict number is :" << number << std::endl; return 0; } 

推理的代码已经编写完成，还需要对其进行编译，根据我们部署的平台，选择编译方式，比如安卓，可以选择交叉编译。这里我们选择部署在本机上。

可以使用 g++进行编译，编译时需要连接 MegEngine Lite 库文件，并且准备好 stb 头文件。

配置环境

由于使用 C++调用 MegEngine Lite 接口，所以我们首先需要编译出 MegEngine Lite 的库。 安装 MegEngine：从源代码编译 MegEngine 。请参考编译 MegEngine Lite。

1. clone MegEngine 工程，进入根目录

git clone --depth=1 [email protected]:MegEngine/MegEngine.git cd MegEngine 

1. 安装 MegEngine 所需的依赖

./third_party/prepare.sh ./third_party/install-mkl.sh 

1. 使用 cmake 进行编译工程得到 c++推理所需的库文件

scripts/cmake-build/host_build.sh 

编译完成后，需要的库文件所在地址为：

MegEngine/build_dir/host/MGE_WITH_CUDA_OFF/MGE_INFERENCE_ONLY_ON/Release/install/lite/

这里为了在 g++编译时添加库文件方便，可以将库文件地址设为环境变量

export LITE_INSTALL_DIR=/path/to/megenginelite-lib #上一步中编译生成的库文件安装路径 export LD_LIBRARY_PATH=$LITE_INSTALL_DIR/lib/x86_64/:$LD_LIBRARY_PATH 

安装 stb： stb 是一个轻量化的图片加载库，可以替代 opencv 完成图片的解码。想要使用它，只需要将对应的头文件包含到项目内，不像 opencv 需要编译产生链接库。

这里为了调用方便直接将 stb 的项目下载下来：

git clone https://github.com/nothings/stb.git 

想要使用图片加载函数 stbi_load(),只需在 cpp 文件中 define STB_IMAGE_IMPLEMENTATION 并且 include stb_image.h 头文件

#define STB_IMAGE_IMPLEMENTATION #include "stb/stb_image.h" 

动态链接编译

最后使用 g++或者 clang 完成对 inference.cpp 的编译。

• -I 选项添加编译时头文件搜索路径
• -l 添加动态链接库
• -L 添加动态链接库搜索路径

g++ -o inference -I$LITE_INSTALL_DIR/include -I./stb inference.cpp -llite_shared -L$LITE_INSTALL_DIR/lib/x86_64 

编译后会在本目录下会得到 inference 二进制文件

执行二进制文件

准备好一张手写数字图片，将图片与模型放到同一目录，执行编译好的文件即可得到推理结果。

./inference test.jpg 

以上就完成了 LeNet 神经网络的部署。

更多 MegEngine 信息获取，您可以： 查看MegEngine 官网和GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705