｜本文转载自 @嘿呀嘿 个人 blog。
1. 前言
Megengine 是旷视科技开发的一款训练推理一体化的深度学习框架，类似于 pytorch，tensorflow。
使用 Megengine 可以快速实现常见的深度学习模型，本文将使用 Megengine 实现手写数字识别，以完成深度学习的两大步骤：训练和预测。通过本文，读者对深度学习的最基本流程和 Megengine 框架的使用方法有大致了解。
2. 环境安装
在命令行输入下列语句即可安装 Megengine ，建议使用 python 版本为 3.5 到 3.8
python3 -m pip install --upgrade pippython3 -m pip install megengine -f https://megengine.org.cn/whl/mge.html
复制代码
安装完成后可以在命令行测试是否安装成功。
python3import megengineprint(megengine.__version__)
复制代码
3. 训练
本部分训练代码来自 Megengine 官方教程，需要详细了解细节请前往 MegEngine 快速上手
3.1 数据集准备
3.1.1 下载数据集
深度学习的第一步为准备数据集，通常会为数据集写一个接口来访问数据集，并对数据进行预处理。
Megengine 中已经实现了 MNIST 数据集的接口，我们可以通过以下代码直接获取。如果想要制作或使用其他数据集，可以点击这里进行学习。
from megengine.data.dataset import MNISTDATA_PATH = "./datasets/MNIST"#第一次运行后，将 download 改为 Falsetrain_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 实现。
功能： 构建可迭代的数据装载器，非常灵活地从数据集连续获取小批量数据参数：
功能：创建一个列表，包含所有数据的索引，可实现数据的随机取样参数：
import megengine.data as dataimport megengine.data.transform as Ttrain_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 Fimport megengine.module as Mclass 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
复制代码
3.3 训练准备
import megengine.optimizer as optimimport megengine.autodiff as autodiffgm = 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 megengineepochs = 10model.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({ "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 LeNetimport cv2import numpy as npimport megengineimport megengine.data.transform as Timport 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):2278–2324, 1998.
附录：train.py
from megengine.data.dataset import MNISTfrom megengine import jit, tensorimport megengineimport numpy as npimport megengine.data as dataimport megengine.data.transform as Timport megengine.functional as Fimport megengine.module as Mimport megengine.optimizer as optimimport 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 modelclass 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）。主要分为三步:
将循环内的前向计算、反向传播和参数优化代码提取成单独的函数，如下面例子中的 train_func()
将网络所需输入作为训练函数的参数，并返回任意你需要的结果（如输出结果、损失函数值等）；
用 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 代码
代码的主逻辑为：
创建 Network
使用 load_model()载入模型
使用 stb 预处理图片（加载和 resize)，然后归一化，载入进 input tensor
使用 network->forward()和 network->wait()完成推理逻辑。
获取模型输出 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。
clone MegEngine 工程，进入根目录
git clone --depth=1 git@github.com:MegEngine/MegEngine.gitcd MegEngine
复制代码
安装 MegEngine 所需的依赖
./third_party/prepare.sh./third_party/install-mkl.sh
复制代码
使用 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 二进制文件
执行二进制文件
准备好一张手写数字图片，将图片与模型放到同一目录，执行编译好的文件即可得到推理结果。
以上就完成了 LeNet 神经网络的部署。
更多 MegEngine 信息获取，您可以：查看文档、和 GitHub 项目，或加入 MegEngine 用户交流 QQ 群：1029741705。欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。
评论