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

作者: Kaami | 来源:发表于2018-07-06 18:58 被阅读31次

    在前不久的某高峰论坛上,小米开源了其移动端的深度学习框架Mobile AI Compute Engine(MACE)。这对于很多致力于嵌入式端优化的人来说,无疑是巨大的惊喜(新坑出现,在 NCNNTVMTensorLite 的坑里不断徘徊的人表示泪目...)。要掌握并使用一个框架,不单单能够跑通 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_partytools文件夹是需要用到的一些第三方模块和代码构建、测试所用到的相关脚本。其中tools/converter.py是构建代码的顶层脚本,有兴趣的可以详细看一下,之前的文章里也有提到 <小米开源框架MACE> 如何构建和使用

    核心的代码在mace文件夹下。同样的,先看一下mace文件夹下的结构:

    仍然是聚焦一些与框架代码相关的文件夹。此处默认读者已事先学习过mace\examples\cli\examples.cc,如没有建议先过一遍,对如何在项目里使用 mace 有个了解。也可以参考 <小米开源框架MACE> 如何构建和使用 一文的最后内容。按照 examples.cc 里的引用,先来看 public 文件夹里的两个头文件:mace.hmace_runtime.h

    mace\public\mace.h 剖析

    mace.h头文件定义了mace框架的几个核心API,包括 CallStatsConvPoolArgsOperatorStats 三个结构体以及 DeviceTypeMaceStatus 两个枚举类型。

    enum DeviceType { CPU = 0, GPU = 2, HEXAGON = 3 }; //框架支持的设备类型及其对应ID.
    
    enum MaceStatus {
      MACE_SUCCESS = 0,
      MACE_INVALID_ARGS = 1,
      MACE_OUT_OF_RESOURCES = 2
    }; // 函数的返回类型
    

    mace.h头文件中还定义了三个类:RunMetadataMaceTensorMaceEngine

    RunMetadata:定义了一个public类型的变量 op_stats,类型为vector <OperatorStats>
    MaceTensor:定义了 MACE 输入输出 tensor,有几个不同形式的构造函数
    MaceEngine:作为 MACE 框架引擎,根据构造函数传入的DeviceType执行相应的功能代码

    MaceTensorMaceEngine两个类中都分别定义了一个嵌套类Impl,这个嵌套类才是真正实现具体功能的类。

    mace\core\mace.cc剖析

    对应于mace\public\mace.h头文件的实现文件为mace\core\mace.cc
    作为深入源码的第一步,下面我们一起看一下这个文件。

    首先是LoadModelDataUnloadModelData两个与模型数据导入和清空的函数。

    接下来是MaceTensor中的嵌套类Impl的定义:

    class MaceTensor::Impl {
     public:
      std::vector<int64_t> shape;
      std::shared_ptr<float> data;
    };
    

    然后是MaceTensor的构造函数实现。MaceTensor的构造函数主要实现Impl类中shapedata两个变量的初始化工作,以下贴出其中一种形式的构造函数,更多代码请查阅mace\core\mace.cc

    MaceTensor::MaceTensor(const std::vector<int64_t> &shape,
                           std::shared_ptr<float> data) {
      MACE_CHECK_NOTNULL(data.get());
      impl_ = std::unique_ptr<MaceTensor::Impl>(new MaceTensor::Impl());
      impl_->shape = shape;
      impl_->data = data;
    }
    

    MaceTensor 类中还有三个成员函数,用于获取impl的变量值。

    const std::vector<int64_t> &MaceTensor::shape() const { return impl_->shape; }
    
    const std::shared_ptr<float> MaceTensor::data() const { return impl_->data; }
    
    std::shared_ptr<float> MaceTensor::data() { return impl_->data; }
    

    接下来是MaceEngine中的嵌套类Impl的定义,其中涉及到的WorkspaceNetBase等类,暂时不需要过多关注,只需要知道这个内嵌类实现了MACE引擎的创建、初始化和运行。下边是Impl类的主体代码:

    class MaceEngine::Impl {
     public:
      explicit Impl(DeviceType device_type);
    
      ~Impl();
    
      MaceStatus Init(const NetDef *net_def,
                      const std::vector<std::string> &input_nodes,
                      const std::vector<std::string> &output_nodes,
                      const unsigned char *model_data);
    
      MaceStatus Init(const NetDef *net_def,
                      const std::vector<std::string> &input_nodes,
                      const std::vector<std::string> &output_nodes,
                      const std::string &model_data_file);
    
      MaceStatus Run(const std::map<std::string, MaceTensor> &inputs,
                     std::map<std::string, MaceTensor> *outputs,
                     RunMetadata *run_metadata);
    
     private:
      const unsigned char *model_data_;
      size_t model_data_size_;
      std::shared_ptr<OperatorRegistry> op_registry_;
      DeviceType device_type_;
      std::unique_ptr<Workspace> ws_;
      std::unique_ptr<NetBase> net_;
      std::map<std::string, mace::InputInfo> input_info_map_;
      std::map<std::string, mace::OutputInfo> output_info_map_;
    #ifdef MACE_ENABLE_HEXAGON
      std::unique_ptr<HexagonControlWrapper> hexagon_controller_;
    #endif
    
      MACE_DISABLE_COPY_AND_ASSIGN(Impl);
    };
    

    Impl 类中有ImplInitRun三个函数(Init有一个重载函数)。

    • Impl函数
      Impl 函数是构造函数,完成一些变量的初始化工作。与 Tensor 相关的一些操作会在 Workspace 下,因而此处创建了一个 Workspace对象。

    • Init函数
      Init 函数完成 MaceEngine 的初始化工作。包括根据网络定义创建输入输出的存储 map,根据输入输出节点和 deviceType 创建Tensor,初始化整个网络模型。Init还有一个重载函数,区别只是const unsigned char *model_dataconst std::string &model_data_file 这两个参数。后者通过 LoadModelData 载入模型数据后调用前者完成初始化。

    • Run函数
      Run 函数根据输入输出文件信息,初始化的模型信息和模型数据,以及 device_type_ 类型调用相应的执行代码。完成模型的前向运算。输出运算结果并保存。

    MaceEngine 类本身还有InitRun 函数,都是通过调用 Impl 类的相应函数实现的。

    mace\core\mace.cc 中还有一个 CreateMaceEngineFromProto 函数。这个函数和 mace/codegen/engine/mace_engine_factory.hCreateMaceEngineFromCode 函数是对应的。两者都是通过调用MaceEngineInit 函数(也即调用 ImplInit函数) 实现初始化过程。 区别在于,前者是通过模型proto定义文件进行初始化,后者是将模型编为代码(在模型部署文件 .yaml 中设置CODE_TYPEcode)。mace/codegen/engine/mace_engine_factory.h 这个文件是代码构建之后产生的。

    以上,对 mace\public\mace.hmace\core\mace.cc 进行了分解,并分析了 mace 框架顶层的几个API,也是框架的基础。

    后续将陆续扩展到框架中核心代码的实现。

    敬请拍砖(轻拍o( ̄▽ ̄)d )。

    版权所有,欢迎转载,转载请注明出处

    相关文章

      网友评论

      • juggzhang:最近也在关注小米的mace,官方提到了模型加密,但是我找不到相关资料,不知道博主有没有了解过,感谢感谢
        Kaami:@juggzhang 就是build时配置文件中model_graph_format和model_data_format的选择。如果是file会生成.pb和.data两个文件,对应模型和数据,部署时就要作为外部加载文件。如果是code,会生成头文件和库文件,模型和数据被编译成.cc文件,这样部署的时候可以一起编译到代码里。显然后一种安全性更好

      本文标题:<小米开源框架MACE> 源码阅读笔记 1

      本文链接:https://www.haomeiwen.com/subject/nnvpuftx.html