减少 MegEngine Lite 体积

midout 是 MegEngine 中用来减小生成的二进制文件体积的工具,有助于在空间受限的设备上部署应用。 midout 通过记录模型推理时用到的 opr 和执行流,使用 if(0) 来关闭未被记录的代码段后重新编译, 重新编译时利用 -flto -ffunction-sections -fdata-sections -Wl,--gc-sections 编译参数,可以大幅度减少目标文件的体积大小(动态库和可执行程序)。 现在基于 MegEngine Lite 提供模型验证工具 Load and Run , 展示怎样在某 AArch64 架构的 Android 端上裁剪 MegEngine 库,Load and Run 底层集成了 MegEngine Lite,所以 使用裁减了 Load and Run 的 Binary size 的大小其实就是对 MegEngine Lite 进行了裁减。

编译静态链接的 load_and_run

编译方法详见 load-and-run 的编译部分。

./cross_build_android_arm_inference.sh

查看一下 load_and_run 的大小:

du ./build_dir/android/arm64-v8a/Release/install/bin/load_and_run
23200

此时 load_and_run 大小超过 20MB. load_and_run 的执行,请参考下文“代码执行”部分。

裁剪 load_and_run

MegEngine 的裁剪可以从两方面进行:

  1. 通过opr 裁剪。在 dump 模型时,可以同时将模型用到的 opr 信息以 json 文件的形式输出, midout 在编译期裁掉没有被模型使用到的所有 opr.

  2. 通过 trace 流裁剪。运行一次模型推理,根据代码的执行流生成 trace 文件, 通过trace文件,在二次编译时将没有执行的代码段裁剪掉。

整个裁剪过程分为两个步骤:

  1. 第一步,dump 模型,获得模型 opr 信息;通过一次推理,获得 trace 文件。

  2. 第二步,使用MegEngine的头文件生成工具 tools/gen_header_for_bin_reduce.py 将 opr 信息和 trace 文件作为输入, 生成 src/bin_reduce_cmake.h CMake 会自动维护这个文件,用户无需关心。 当然也可以单独使用模型 opr 信息或是 trace 文件来生成 src/bin_reduce_cmake.h , 单独使用 opr 信息时,默认保留所有 kernel,单独使用 trace 文件时,默认保留所有 opr.

dump 模型获得 opr 类型名称

一个模型通常不会用到所有的 opr,根据模型使用的 opr,可以裁掉那些模型没有使用的 opr。 在模型 dump 时,我们可以获得模型的 opr 信息。 使用 dump 模型时候加上 strip_info_file 参数,可以将模型使用的 opr 信息 dump 到一个 JSON 文件中,如果 需要将多个模型的 opr JSON 文件拼接在一起,可以在 dump 模型时候再加上 append_json 参数。

import numpy as np
import megengine.functional as F
import megengine.hub
from megengine import jit, tensor

if __name__ == "__main__":
   net = megengine.hub.load("megengine/models", "resnet50", pretrained=True)
   net.eval()

   @jit.trace(symbolic=True, capture_as_const=True)
   def fun(data, *, net):
      pred = net(data)
      pred_normalized = F.softmax(pred)
      return pred_normalized

   data = tensor(np.random.random([1, 3, 224, 224]).astype(np.float32))

   fun(data, net=net)
   fun.dump("resnet50.mge", arg_names=["data"], optimize_for_inference=True, strip_info_file = "resnet50.mge.json")

执行完毕后,会生成 resnet50.mgeresnet50.mge.json . 查看这个 JSON 文件,它记录了模型用到的 opr 名称。

cat resnet50.mge.json
{"hash": 238912597679531219, "dtypes": ["Byte", "Float32", "Int32"], "opr_types": ["Concat", "ConvBiasForward", "ConvolutionForward", "Elemwise", "GetVarShape", "Host2DeviceCopy", "ImmutableTensor", "MatrixMul", "MultipleDeviceTensorHolder", "PoolingForward", "Reshape", "Subtensor"], "elemwise_modes": ["ADD", "FUSE_ADD_RELU"]}

执行模型获得 trace 文件

基于 trace 的裁剪需要通过一次推理获得模型的执行 trace 文件。具体步骤如下:

  1. CMake 构建时,打开 MGE_WITH_MIDOUT_PROFILE 开关,编译 load_and_run:

    EXTRA_CMAKE_ARGS="-DMGE_WITH_MIDOUT_PROFILE=ON" ./cross_build_android_arm_inference.sh -r
    

    编译完成后,将 build_dir/android/arm64-v8a/Release/install/bin 下的 load_and_run 推至设备并执行:

    ./load_and_run ./resnet50.mge
    

    得到如下输出(x.x.x 指代当前版本信息,下同):

    mgb load-and-run: using MegEngine MegEngine x.x.x and MegDNN x.x.x
    load model: 70.888ms
    === going to run 1 testcases; output vars: ADD(reshape[2655],reshape[2663])[2665]{1,1000}
    === prepare: 4.873ms; going to warmup
    warmup 0: 877.578ms
    === going to run test #0 for 10 times
    iter 0/10: 481.445ms (exec=481.436,device=480.794)
    iter 1/10: 481.192ms (exec=481.183,device=481.152)
    iter 2/10: 480.430ms (exec=480.420,device=480.389)
    iter 3/10: 479.593ms (exec=479.585,device=479.553)
    iter 4/10: 479.851ms (exec=479.843,device=479.811)
    iter 5/10: 479.581ms (exec=479.572,device=479.541)
    iter 6/10: 480.174ms (exec=480.165,device=480.134)
    iter 7/10: 479.443ms (exec=479.435,device=479.404)
    iter 8/10: 479.987ms (exec=479.978,device=479.948)
    iter 9/10: 480.637ms (exec=480.628,device=480.598)
    === finished test #0: time=4802.333ms avg_time=480.233ms sd=0.688ms minmax=479.443,481.445
    
    === total time: 4802.333ms
    midout: 110 items written to midout_trace.20717
    

    注意到执行模型后,生成了 midout_trace.20717 文件(20717 是进程的 PID,每次运行可能不一样),该文件记录了模型在底层执行了哪些 kernel。

  2. 生成 src/bin_reduce_cmake.h 并再次编译 load_and_run:

    将生成的 midout_trace.20717 拷贝至本地, 使用上文提到的头文件生成工具 gen_header_for_bin_reduce.py 生成 src/bin_reduce_cmake.h .

    python3 ./tools/gen_header_for_bin_reduce.py resnet50.mge.json midout_trace.20717 -o src/bin_reduce_cmake.h
    
    EXTRA_CMAKE_ARGS="-DMGE_WITH_MINIMUM_SIZE=ON" ./scripts/cmake-build/cross_build_android_arm_inference.sh -r
    

    编译完成后,检查 load_and_run 的大小, 注意 MGE_WITH_MINIMUM_SIZE 不是非必须的,加上它 size 会更小,但同时会关闭一些编译选项(详情可参考 MegEngine 工程根目录的 CMakeLists.txt):

    du build_dir/android/arm64-v8a/release/install/bin/load_and_run
    2264
    

    此时 load_and_run 的大小减小到 2MB 多。推到设备上运行,得到如下输出:

    mgb load-and-run: using MegEngine x.x.x and MegDNN x.x.x
    load model: 74.208ms
    === going to run 1 testcases; output vars: ADD(reshape[2655],reshape[2663])[2665]{1,1000}
    === prepare: 1.251ms; going to warmup
    warmup 0: 377.813ms
    === going to run test #0 for 10 times
    iter 0/10: 266.996ms (exec=266.993,device=266.854)
    iter 1/10: 266.717ms (exec=266.715,device=266.702)
    iter 2/10: 266.867ms (exec=266.865,device=266.855)
    iter 3/10: 267.172ms (exec=267.171,device=267.159)
    iter 4/10: 266.820ms (exec=266.819,device=266.807)
    iter 5/10: 266.852ms (exec=266.850,device=266.838)
    iter 6/10: 267.376ms (exec=267.374,device=267.363)
    iter 7/10: 267.005ms (exec=267.003,device=266.991)
    iter 8/10: 266.685ms (exec=266.684,device=266.671)
    iter 9/10: 266.767ms (exec=266.766,device=266.755)
    === finished test #0: time=2669.257ms avg_time=266.926ms sd=0.216ms minmax=266.685,267.376
    
    === total time: 2669.257ms
    

可以看到模型依然正常运行,并且运行速度正常。

使用裁剪后的 load_and_run

想要裁剪前后的应用能够正常运行,需要保证裁剪前后两次推理使用同样的命令行参数以及相同的 CPU 平台。 如果使用上文裁剪的 load_and_fun 的 fast-run功能(详见 使用 Load and run 测试与验证模型 )。

./load_and_run resnet50.mge --fast-run --fast-run-algo-policy resnet50.cache

可能得到如下输出:

mgb load-and-run: using MegEngine x.x.x and MegDNN x.x.x
load model: 71.927ms
=== going to run 1 testcases; output vars: ADD(reshape[2655],reshape[2663])[2665]{1,1000}
=== prepare: 1.251ms; going to warmup
 Trap

这是因为程序运行到了已经被裁剪掉的函数中,未被记录在 trace 文件中的函数的实现已经被替换成 trap() . 如果想要裁剪与 fast-run 配合使用,需要按如下流程获得 trace 文件:

  1. 开启 fast-run 模式,执行未裁剪的 load_and_run 获得 .cache 文件,注意本次执行生成的 trace 应该被丢弃:

    ./load_and_run resnet50.mge --fast-run --fast-run-algo-policy resnet50.cache
    
  2. 使用 .cache 文件,执行 load_and_run 获得 trace 文件:

    ./load_and_run resnet50.mge --fast-run-algo-policy resnet50.cache
    
  3. 如上节,将 trace 文件拷贝回本机,生成 src/bin_reduce_cmake.h ,再次编译 load_and_run 并推至设备。

  4. 使用裁剪后的 load_and_run 的 fast-run 功能,执行同 2 的命令,得到如下输出:

    mgb load-and-run: using MegEngine x.x.x and MegDNN x.x.x
    [04 15:34:18 from_argv@mgblar.cpp:1392][WARN] enable winograd transform
    load model: 64.228ms
    === going to run 1 testcases; output vars: ADD(reshape[2655],reshape[2663])[2665]{1,1000}
    === prepare: 260.058ms; going to warmup
    warmup 0: 279.550ms
    === going to run test #0 for 10 times
    iter 0/10: 209.177ms (exec=209.164,device=209.031)
    iter 1/10: 209.010ms (exec=209.008,device=208.997)
    iter 2/10: 209.024ms (exec=209.022,device=209.011)
    iter 3/10: 208.584ms (exec=208.583,device=208.573)
    iter 4/10: 208.669ms (exec=208.667,device=208.658)
    iter 5/10: 208.849ms (exec=208.847,device=208.838)
    iter 6/10: 208.787ms (exec=208.785,device=208.774)
    iter 7/10: 208.703ms (exec=208.701,device=208.692)
    iter 8/10: 208.918ms (exec=208.916,device=208.905)
    iter 9/10: 208.669ms (exec=208.667,device=208.656)
    === finished test #0: time=2088.390ms avg_time=208.839ms sd=0.191ms minmax=208.584,209.177
    
    === total time: 2088.390ms
    

使用其他 load_and_run 提供的功能也是如此,想要裁剪前后的应用能够正常运行, 需要保证裁剪前后两次推理使用同样的命令行参数。

多个模型合并裁剪

多个模型的合并裁剪与单个模型流程相同。 gen_header_for_bin_reduce.py 接受多个输入。 假设有模型 A 与模型 B, 已经获得 A.mge.json , B.mge.json 以及 A.trace , B.trace . 执行:

python3 ./tools/gen_header_for_bin_reduce.py A.mge.json A.trace B.mge.json B.trace -o src/bin_reduce_cmake.h

裁剪基于 MegEngine Lite 的应用

可以通过如下几种方式集成 MegEngine Lite,对应的裁剪方法相差无几:

  1. 参照 MegEngine Lite C++ 模型部署快速上手 ,将整个 MegEngine Lite 集成到用户工程中。 只需要按照上文中裁剪 load_and_run 的流程裁剪用户的工程即可。

  2. 可能一个应用想要通过静态库集成 MegEngine Lite。此时需要获得一个裁剪过的 liblite_static_all_in_one.a . 可以依然使用 load_and_run 运行模型获得 trace 文件, 生成 bin_reduce.h ,并二次编译获得裁剪过的 liblite_static_all_in_one.a . 此时,用户使用自己编写的构建脚本构建应用程序,并静态链接 liblite_static_all_in_one.a , 加上链接参数 -flto=full -ffunction-sections -fdata-sections -Wl,--gc-sections . 即可得到裁剪过的基于 MegEngine 的应用。

  3. 上述流程亦可以用于 liblite_shared.so 的裁剪,当使用动态库进行裁剪时, 用户自己的编译脚本不再强求添加 -flto=full -ffunction-sections -fdata-sections -Wl,--gc-sections 的编译选项。

警告

midout 裁减之后的应用程序有一些限制:

  • 模型的输入数据的 shape 不能改变

  • 模型必须是静态图模型,不是有动态分支

  • 裁减之后的应用程序只能在裁减时候的平台上运行,其他平台改变了程序的运行路径可能会失败, 比如 CPU 架构的不同。

注解

如果裁减的应用程序需要在不同的 shape 下面都能成功运行,需要在裁减 执行模型获得 trace 文件 中所有的输入都运行一次。