MegEngine 作為一個(gè)訓(xùn)推一體的 AI 框架，為用戶提供了模型訓(xùn)練以及部署的能力。但是在部署模型時(shí)，由于會(huì)存在對(duì)于部署的模型計(jì)算來(lái)說(shuō)不必要的代碼，導(dǎo)致 SDK 體積比較大。為了解決上述問(wèn)題，我們提供了新的工具：AI 編譯器 MegCC。

MegCC 有以下特性：

1. 只生成完成輸入模型計(jì)算的必要的 kernel 代碼以及用戶指定的 CV 算子，編譯得到的二進(jìn)制文件體積很小。

2. 支持 int8 量化模型以及 float16 量化，且生成的 kernel 是精心優(yōu)化過(guò)的，推理性能好。

3. 支持平臺(tái)廣。硬件方面，支持 Armv8、Armv7 以及 X86。操作系統(tǒng)方面，支持標(biāo)準(zhǔn)和非標(biāo)準(zhǔn)操作系統(tǒng)。

本文將重點(diǎn)解析模型部署中的重要步驟之一-模型編譯：編譯 MegEngine 模型，生成運(yùn)行這個(gè)模型對(duì)應(yīng)的 Kernel 以及和這些 Kernel 綁定的模型。

編譯模型時(shí)：

• MegCC 生成模型使用的內(nèi)核和用戶所需的 CV 內(nèi)核
• MegCC 做了多項(xiàng)優(yōu)化，例如靜態(tài)內(nèi)存規(guī)劃和模型優(yōu)化
• MegCC 將上述數(shù)據(jù)轉(zhuǎn)儲(chǔ)到最終模型中

模型編譯階段主要使用 mgb-to-tinynn 工具，編譯完成之后，會(huì)在用戶給定的目錄下面，生成對(duì)應(yīng)的純 C 代碼的 Kernel 以及對(duì)應(yīng)的模型。為了編譯模型，mgb-to-tinynn 工具需要用戶提供一個(gè) Json 文件來(lái)配置編譯選項(xiàng)。

目前 MegCC 只支持 mge 模型作為輸入，其他模型格式可以考慮轉(zhuǎn)換到 ONNX，然后通過(guò) mgeconvert 進(jìn)行模型格式轉(zhuǎn)換。

編寫(xiě) Json 文件

json模板如下：

{
    "@dump_dir":"[Required], specify the directory where the output kernel and model are stored",
    "dump_dir":"./batch_dump/",
    "models":[
        {
            "@model_name":"[Optional], specify the name of the tiny model to be generated",
            "model_name":"det_nchw44",
            "@model_path":"[Required], specify the input model path. `mge' and `emod' formats are supported.",
            "model_path":"path/to/model.mge",
            "@input_shape_str":"[Optional], modify the input shape",
            "input_shape_str":"data=(1,1,384,288):data=(1,1,288,384)",
            "@enable_nchw44":"[Optional], whether to enable nchw44 optimization, default false",
            "enable_nchw44":true,
            "@enable_nchw44_dot":"[Optional], whether to enable nchw44 dot optimization for int8, default false",
            "enable_nchw44_dot":false,
            "@add_nhwc2nchw_to_input":"[Optional], add nhwc2nchw dimshuffle to input",
            "add_nhwc2nchw_to_input":false,
            "@mgb_fuse_kernel":"[Optional], fuse mgb kernel as possible",
            "mgb_fuse_kernel":false,
            "@enable_compress_fp16":"[Optional], whether to enable the optimization of using float16 storage to compress the model size",
            "enable_compress_fp16":false,
            "@enable_nchw88":"[Optional], whether to enable nchw88 optimization, default false",
            "enable_nchw88":false,
            "@enable_ioc16":"[Optional], whether to enable optimization using float16 calculation, default false",
            "enable_ioc16":false
        },
        {
            "model_name":"pf_nchw44",
            "model_path":"path/to/another_model.emod",
            "input_shape_str":"data=(1,1,112,112)",
            "enable_nchw44":true
        }
    ],
    "@cv":"[Optional], specify the cv operator used in non-models (e.g. in pre and post processing)",
    "cv":{
        "transpose":["ui8"],
        "roicopy":["ui8"],
        "rotate":["ui8"],
        "flip":["ui8"],
        "resize_linear":["ui8", "f32"],
        "warp_affine_replicate_linear":["ui8"],
        "rgb2bgr":["ui8"],
        "yuv2bgr_nv21":["ui8"],
        "rgb2yuv":["ui8"]
    }
}

• 設(shè)置模型編譯之后 dump 的路徑，可以在 mgb-to-tinynn 工具中通過(guò) --dump 參數(shù)進(jìn)行 override。

• Json 文件中需要指定使用 mgb-to-tinynn 編譯的模型名稱，模型的路徑，以及模型的輸入數(shù)據(jù)，以及一些優(yōu)化參數(shù)等

  • 如果部署的實(shí)際情況中需要多個(gè)模型組成 pipline，可以指定多個(gè)模型
  • 如果一個(gè)模型在實(shí)際推理過(guò)程中可能需要多種輸入 shape，需要分別在 input_shape_str 中指定，并用 : 分割開(kāi)。
  • 支持 enable_nchw44 和 enable_nchw44_dot 兩個(gè)優(yōu)化選項(xiàng)，enable_nchw44 為 true 表示，優(yōu)化模型推理中 Tensor layout 為 NC4HW4。enable_nchw44_dot 為 true 表示，優(yōu)化模型推理中 Tensor layout 為 NC4HW4，并且在推理過(guò)程中使用 ArmV8.2 dot 指令進(jìn)行推理加速
  • 開(kāi)啟 enable_ioc16 優(yōu)化選項(xiàng)，可以使用 float16 進(jìn)行 float32 模型的計(jì)算，可提升推理性能，代價(jià)是損失一些精度。通常需要同時(shí)開(kāi)啟 enable_nchw88 選項(xiàng)

• 另外為了方便用戶集成時(shí)候使用 cv 算子進(jìn)行模型的前后處理，可以在這個(gè) Json 文件中指定需要用到的 cv 算子的名稱以及對(duì)應(yīng)的數(shù)據(jù)類型。MegCC 支持的 cv 算子 列表。

模型編譯

編譯模型目前可以使用 mgb-to-tinynn 這個(gè)可執(zhí)行文件完成編譯，也可以使用 Release 包里面的現(xiàn)成腳本 ./script/ppl_gen.sh 進(jìn)行編譯。

使用現(xiàn)成腳本進(jìn)行模型編譯（推薦）

Release 包中的 script 目錄下面有一個(gè) ppl_gen.sh 的文件，直接執(zhí)行：

./script/ppl_gen.sh ./bin/mgb-to-tinynn ./example/mobilenet.json mobilenet_gen --arm64

./script/ppl_gen.sh 這個(gè)腳本將執(zhí)行模型編譯，并把 Runtime 需要的資源一同打包在一個(gè)壓縮包中，方便后續(xù) Runtime 的編譯，解壓這個(gè)壓縮包將得到：

├── build runtime build 的路徑
├── immigration generalIntrinsic 頭文件
│   └── include
├── kern 模型 kernel 文件包括cv 算子
├── mobilenet.json 模型dump所用的配置文件
├── model 模型
│   └── mobilenet_nchw44.tiny
├── model_info 模型輸入信息
│   └── mobilenet_nchw44.tiny.txt
├── ppl_build.sh
├── runtime runtime 源碼
│   ├── CMakeLists.txt
│   ├── example
│   ├── flatcc
│   ├── include
│   ├── schema
│   ├── script
│   └── src
└── test_model.py 模型對(duì)分測(cè)試腳本

使用可執(zhí)行文件編譯

使用 mgb-to-tinynn 和上面寫(xiě)好的 Json 文件執(zhí)行：

mgb-to-tinynn --json=/path/to/json --[target]

成模型編譯后，將生成的運(yùn)行這個(gè)模型的 Kernel，和這些 Kernel 綁定的模型文件以及 cv 算子都放在 Json 文件中指定的目錄。其中

• target：可以是 baremetal, arm64, armv7, arm64v7.

• baremetal: 生成的 Kernel 為單片機(jī)可以運(yùn)行的純 C 形式
• arm64v7: 生成能夠同時(shí)在 Arm64 和 ArmV7 上運(yùn)行的兩套 Kernel 以及他們對(duì)應(yīng)的模型，這時(shí)候，模型文件可能會(huì)比 target 為 arm64 時(shí)候大一點(diǎn)。

如編譯 Release 包中的 mobilenet 模型，目標(biāo)機(jī)器是 arm64 機(jī)器，運(yùn)行如下命令

mkdir mobilenet_gen
./bin/mgb-to-tinynn --json=./example/mobilenet.json --arm64 --dump mobilenet_gen