DSP

小米开源框架MACE 源码阅读笔记

2019-07-13 15:26发布

站内文章 / DSP
15781 0 1580



转载自 https://www.jianshu.com/p/7061fd67d419

前扯

在前不久的某高峰论坛上，小米开源了其移动端的深度学习框架Mobile AI Compute Engine（MACE）。这对于很多致力于嵌入式端优化的人来说，无疑是巨大的惊喜（新坑出现，在 NCNN、TVM、TensorLite 的坑里不断徘徊的人表示泪目…）。要掌握并使用一个框架，不单单能够跑通 demo，还需要对其源码有充分的认识，知其优而优。笔者开启此帖（坑），希望能够分享自己学习源码的过程，水平及理解有限，欢迎交流（点赞）和指正（赞赏）。另外，Github上的代码会一直处于更新中，后续会尽量紧跟更新。

开篇

首先看一下从Github上拉下来的代码结构：下边的几个文件是一些版本发布与介绍性文件，与代码的相关性不大，简单看下 README_zh.md 里边说的：
“Mobile AI Compute Engine (MACE) 是一个专为移动端异构计算平台优化的神经网络计算框架。主要从以下的角度做了专门的优化：
  • 性能
    代码经过NEON指令，OpenCL以及Hexagon HVX专门优化，并且采用 Winograd算法 来进行卷积操作的加速。此外，还对启动速度进行了专门的优化。
  • 功耗
    支持芯片的功耗管理，例如ARM的big.LITTLE调度，以及高通Adreno GPU功耗选项。
  • 系统响应
    支持自动拆解长时间的OpenCL计算任务，来保证UI渲染任务能够做到较好的抢占调度，从而保证系统UI的相应和用户体验。
  • 内存占用
    通过运用内存依赖分析技术，以及内存复用，减少内存的占用。另外，保持尽量少的外部依赖，保证代码尺寸精简。
  • 模型加密与保护
    模型保护是重要设计目标之一。支持将模型转换成C++代码，以及关键常量字符混淆，增加逆向的难度。
  • 硬件支持范围
    支持高通，联发科，以及松果等系列芯片的CPU，GPU与DSP(目前仅支持Hexagon)计算加速。同时支持在具有POSIX接口的系统的CPU上运行。
    ”
看其介绍，MACE团队应该是倾注了不少心血，希望能够不断完善下去。回到代码目录上来。docker文件夹是几个与docker安装相关的文件，不用过多关注。docs文件夹是一些文档，与代码关系不大。third_party和tools文件夹是需要用到的一些第三方模块和代码构建、测试所用到的相关脚本。其中tools/converter.py是构建代码的顶层脚本，有兴趣的可以详细看一下，之前的文章里也有提到 <小米开源框架MACE> 如何构建和使用。核心的代码在mace文件夹下。同样的，先看一下mace文件夹下的结构：仍然是聚焦一些与框架代码相关的文件夹。此处默认读者已事先学习过maceexamplescliexamples.cc，如没有建议先过一遍，对如何在项目里使用 mace 有个了解。也可以参考 <小米开源框架MACE> 如何构建和使用 一文的最后内容。按照 examples.cc 里的引用，先来看 public 文件夹里的两个头文件：mace.h和mace_runtime.h。

macepublicmace.h 剖析

mace.h头文件定义了mace框架的几个核心API，包括 CallStats、ConvPoolArgs、OperatorStats 三个结构体以及 DeviceType 和 MaceStatus 两个枚举类型。
  1. enum DeviceType { CPU = 0, GPU = 2, HEXAGON = 3 }; //框架支持的设备类型及其对应ID.
  1. enum MaceStatus {
  2. MACE_SUCCESS = 0,
  3. MACE_INVALID_ARGS = 1,
  4. MACE_OUT_OF_RESOURCES = 2
  5. }; // 函数的返回类型
mace.h头文件中还定义了三个类：RunMetadata、MaceTensor 和 MaceEngine
RunMetadata：定义了一个public类型的变量 op_stats，类型为vector
MaceTensor：定义了 MACE 输入输出 tensor，有几个不同形式的构造函数
MaceEngine：作为 MACE 框架引擎，根据构造函数传入的DeviceType执行相应的功能代码
MaceTensor 和 MaceEngine两个类中都分别定义了一个嵌套类Impl，这个嵌套类才是真正实现具体功能的类。

macecoremace.cc剖析

对应于macepublicmace.h头文件的实现文件为macecoremace.cc
作为深入源码的第一步，下面我们一起看一下这个文件。首先是LoadModelData 和 UnloadModelData两个与模型数据导入和清空的函数。接下来是MaceTensor中的嵌套类Impl的定义：
  1. class MaceTensor::Impl {
  2. public:
  3. std::vector<int64_t> shape;
  4. std::shared_ptr<float> data;
  5. };
然后是MaceTensor的构造函数实现。MaceTensor的构造函数主要实现Impl类中shape和data两个变量的初始化工作，以下贴出其中一种形式的构造函数，更多代码请查阅macecoremace.cc
  1. MaceTensor::MaceTensor(const std::vector<int64_t> &shape,
  2. std::shared_ptr<float> data) {
  3. MACE_CHECK_NOTNULL(data.get());
  4. impl_ = std::unique_ptr(new MaceTensor::Impl());
  5. impl_->shape = shape;
  6. impl_->data = data;
  7. }
MaceTensor 类中还有三个成员函数，用于获取impl的变量值。
  1. const std::vector<int64_t> &MaceTensor::shape() const { return impl_->shape; }
  3. const std::shared_ptr<float> MaceTensor::data() const { return impl_->data; }
  5. std::shared_ptr<float> MaceTensor::data() { return impl_->data; }
接下来是MaceEngine中的嵌套类Impl的定义，其中涉及到的Workspace、NetBase等类，暂时不需要过多关注，只需要知道这个内嵌类实现了MACE引擎的创建、初始化和运行。下边是Impl类的主体代码：
  1. class MaceEngine::Impl {
  2. public:
  3. explicit Impl(DeviceType device_type);
  5. ~Impl();
  7. MaceStatus Init(const NetDef *net_def,
  8. const std::vector<std::string> &input_nodes,
  9. const std::vector<std::string> &output_nodes,
  10. const unsigned char *model_data);
  12. MaceStatus Init(const NetDef *net_def,
  13. const std::vector<std::string> &input_nodes,
  14. const std::vector<std::string> &output_nodes,
  15. const std::string &model_data_file);
  17. MaceStatus Run(const std::map<std::string, MaceTensor> &inputs,
  18. std::map<std::string, MaceTensor> *outputs,
  19. RunMetadata *run_metadata);
  21. private:
  22. const unsigned char *model_data_;
  23. size_t model_data_size_;
  24. std::shared_ptr op_registry_;
  25. DeviceType device_type_;
  26. std::unique_ptr ws_;
  27. std::unique_ptr net_;
  28. std::map<std::string, mace::InputInfo> input_info_map_;
  29. std::map<std::string, mace::OutputInfo> output_info_map_;
  30. #ifdef MACE_ENABLE_HEXAGON
  31. std::unique_ptr hexagon_controller_;
  32. #endif
  34. MACE_DISABLE_COPY_AND_ASSIGN(Impl);
  35. };
Impl 类中有Impl、Init、Run三个函数（Init有一个重载函数）。
  • Impl函数
    Impl 函数是构造函数，完成一些变量的初始化工作。与 Tensor 相关的一些操作会在 Workspace 下，因而此处创建了一个 Workspace对象。
  • Init函数
    Init 函数完成 MaceEngine 的初始化工作。包括根据网络定义创建输入输出的存储 map，根据输入输出节点和 deviceType 创建Tensor，初始化整个网络模型。Init还有一个重载函数，区别只是const unsigned char *model_data 和 const std::string &model_data_file 这两个参数。后者通过 LoadModelData 载入模型数据后调用前者完成初始化。
  • Run函数
    Run 函数根据输入输出文件信息，初始化的模型信息和模型数据，以及 device_type_类型调用相应的执行代码。完成模型的前向运算。输出运算结果并保存。
MaceEngine 类本身还有Init、Run 函数，都是通过调用 Impl 类的相应函数实现的。macecoremace.cc 中还有一个 CreateMaceEngineFromProto 函数。这个函数和 mace/codegen/engine/mace_engine_factory.h 中 CreateMaceEngineFromCode 函数是对应的。两者都是通过调用MaceEngine 的 Init 函数（也即调用 Impl 的Init函数） 实现初始化过程。 区别在于，前者是通过模型proto定义文件进行初始化，后者是将模型编为代码（在模型部署文件 .yaml 中设置CODE_TYPE 为 code）。mace/codegen/engine/mace_engine_factory.h 这个文件是代码构建之后产生的。以上，对 macepublicmace.h 和 macecoremace.cc 进行了分解，并分析了 mace 框架顶层的几个API，也是框架的基础。后续将陆续扩展到框架中核心代码的实现。敬请拍砖（轻拍o(￣▽￣)ｄ ）。
版权所有，欢迎转载，转载请注明出处

热门文章