From 4f5d36ec2182bf0ad8c58d88b996f5646be5ac70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=A4=9A=E6=AC=A3?= Date: Tue, 28 Sep 2021 14:02:50 +0800 Subject: [PATCH] merge tmconvert --- .github/workflows/ci.yml | 30 - .gitignore | 5 +- README.md | 150 +- bin/convert | 273 ++-- ci/build_cambricon.sh | 12 - ci/pytest_caffe_and_onnx.sh | 39 +- ci/pytest_cambricon.sh | 23 - ci/pytest_tflite.sh | 32 +- ci/pytest_transform.sh | 21 - ci/utils.sh | 5 - mgeconvert/__init__.py | 1 - mgeconvert/backend/__init__.py | 0 .../ir_to_caffe}/__init__.py | 3 +- .../backend/ir_to_caffe/caffe_converter.py | 117 ++ .../ir_to_caffe}/caffe_op.py | 978 ++++++++----- .../ir_to_caffe}/init.sh | 2 +- mgeconvert/backend/ir_to_onnx/__init__.py | 1 + .../ir_to_onnx}/init.sh | 0 .../ir_to_onnx}/onnx_converter.py | 66 +- mgeconvert/backend/ir_to_onnx/onnx_op.py | 1245 ++++++++++++++++ mgeconvert/backend/ir_to_tflite/__init__.py | 9 + mgeconvert/backend/ir_to_tflite/init.sh | 68 + .../ir_to_tflite}/pyflexbuffers/__init__.py | 0 .../pyflexbuffers/fbconverter.cc | 0 .../ir_to_tflite}/tflite_converter.py | 191 +-- .../ir_to_tflite}/tflite_op.py | 618 ++++---- mgeconvert/caffe_converter/caffe_converter.py | 151 -- mgeconvert/cambricon_converter/.gitignore | 3 - mgeconvert/cambricon_converter/CMakeLists.txt | 42 - mgeconvert/cambricon_converter/README.md | 29 - .../cambricon_converter.py | 235 --- .../cambricon_converter/cambricon_op.py | 340 ----- .../cambricon_converter/cmake/cndev.cmake | 48 - .../cambricon_converter/cmake/cnml.cmake | 44 - .../cambricon_converter/cmake/cnrt.cmake | 44 - mgeconvert/cambricon_converter/init.sh | 11 - mgeconvert/cambricon_converter/lib/model.py | 23 - .../cambricon_converter/lib/operators.py | 1265 ----------------- mgeconvert/cambricon_converter/lib/tensor.py | 365 ----- .../cambricon_converter/swig/cambricon.i | 529 ------- .../{utils => converter_ir}/__init__.py | 1 - mgeconvert/converter_ir/ir_graph.py | 157 ++ mgeconvert/converter_ir/ir_op.py | 401 ++++++ mgeconvert/converter_ir/ir_quantizer.py | 86 ++ mgeconvert/converter_ir/ir_tensor.py | 141 ++ mgeconvert/converter_ir/ir_transform.py | 890 ++++++++++++ .../__init__.py | 2 - mgeconvert/converters/mge_to_caffe.py | 36 + mgeconvert/converters/mge_to_onnx.py | 49 + mgeconvert/converters/mge_to_tflite.py | 58 + mgeconvert/converters/tm_to_caffe.py | 57 + mgeconvert/converters/tm_to_onnx.py | 54 + mgeconvert/converters/tm_to_tflite.py | 110 ++ .../__init__.py | 2 - mgeconvert/frontend/mge_to_ir/__init__.py | 8 + mgeconvert/frontend/mge_to_ir/mge_frontend.py | 65 + .../mge_to_ir}/mge_utils.py | 3 +- .../mge_to_ir/op_generators}/__init__.py | 6 +- .../frontend/mge_to_ir/op_generators/base.py | 62 + .../mge_to_ir/op_generators/convolution.py | 174 +++ .../mge_to_ir/op_generators/elemwise.py | 71 + .../mge_to_ir/op_generators/tensor.py | 221 +++ .../frontend/mge_to_ir/symbolvar_resolver.py | 64 + mgeconvert/frontend/tm_to_ir/__init__.py | 8 + .../tm_to_ir/op_generators/__init__.py | 29 + .../frontend/tm_to_ir/op_generators/base.py | 50 + .../tm_to_ir/op_generators/batchnorm.py | 88 ++ .../tm_to_ir/op_generators/broadcast.py | 39 + .../frontend/tm_to_ir/op_generators/concat.py | 80 ++ .../tm_to_ir/op_generators/constant.py | 22 + .../frontend/tm_to_ir/op_generators/conv2d.py | 160 +++ .../frontend/tm_to_ir/op_generators/deconv.py | 35 + .../tm_to_ir/op_generators/dropout.py | 44 + .../tm_to_ir/op_generators/elemwise.py | 288 ++++ .../tm_to_ir/op_generators/flatten.py | 39 + .../tm_to_ir/op_generators/getvarshape.py | 26 + .../frontend/tm_to_ir/op_generators/matmul.py | 134 ++ .../tm_to_ir/op_generators/pooling.py | 89 ++ .../frontend/tm_to_ir/op_generators/reduce.py | 42 + .../tm_to_ir/op_generators/reshape.py | 61 + .../frontend/tm_to_ir/op_generators/resize.py | 45 + .../tm_to_ir/op_generators/softmax.py | 82 ++ .../tm_to_ir/op_generators/squeeze.py | 38 + .../tm_to_ir/op_generators/subtensor.py | 77 + .../tm_to_ir/op_generators/transpose.py | 33 + .../tm_to_ir/op_generators/typecvt.py | 33 + mgeconvert/frontend/tm_to_ir/pattern_utils.py | 161 +++ mgeconvert/frontend/tm_to_ir/qat_pattern.py | 186 +++ mgeconvert/frontend/tm_to_ir/tm_frontend.py | 154 ++ .../frontend/tm_to_ir/tm_tensor_resolver.py | 86 ++ mgeconvert/frontend/tm_to_ir/tm_utils.py | 12 + mgeconvert/mge_context/__init__.py | 16 - mgeconvert/mge_context/mge_net.py | 144 -- mgeconvert/mge_context/mge_op.py | 464 ------ mgeconvert/mge_context/mge_tensor.py | 68 - mgeconvert/mge_context/mge_transform.py | 959 ------------- mgeconvert/onnx_converter/onnx_op.py | 752 ---------- mgeconvert/tflite_converter/__init__.py | 9 - mgeconvert/tflite_converter/init.sh | 41 - mgeconvert/utils/convert_caffe.py | 47 - mgeconvert/utils/convert_cambricon.py | 44 - mgeconvert/utils/convert_onnx.py | 45 - mgeconvert/utils/convert_tflite.py | 47 - mgeconvert/version.py | 5 +- setup.py | 38 +- test/__init__.py | 1 - test/convolution-backward-filter.mge | 3 - test/mge/__init__.py | 7 + test/mge/convolution-backward-filter.mge | Bin 0 -> 1324 bytes test/{ => mge}/test_caffe.py | 20 +- test/{ => mge}/test_onnx.py | 75 +- test/{ => mge}/test_tflite.py | 192 +-- test/quantization_utils.py | 102 -- test/test_cambricon.py | 151 -- test/test_transform.py | 96 -- test/traced_module/__init__.py | 7 + test/traced_module/test_caffe.py | 249 ++++ test/traced_module/test_onnx.py | 228 +++ test/traced_module/test_qat_tflite.py | 165 +++ test/traced_module/test_tflite.py | 240 ++++ .../traced_module/tm_utils.py | 13 +- test/utils.py | 92 +- tools/change_batch.py | 7 + tools/convert_and_check.py | 7 + 124 files changed, 8809 insertions(+), 7402 deletions(-) delete mode 100755 ci/build_cambricon.sh delete mode 100755 ci/pytest_cambricon.sh delete mode 100755 ci/pytest_transform.sh delete mode 100644 ci/utils.sh create mode 100644 mgeconvert/backend/__init__.py rename mgeconvert/{onnx_converter => backend/ir_to_caffe}/__init__.py (83%) create mode 100644 mgeconvert/backend/ir_to_caffe/caffe_converter.py rename mgeconvert/{caffe_converter => backend/ir_to_caffe}/caffe_op.py (51%) rename mgeconvert/{caffe_converter => backend/ir_to_caffe}/init.sh (88%) create mode 100644 mgeconvert/backend/ir_to_onnx/__init__.py rename mgeconvert/{onnx_converter => backend/ir_to_onnx}/init.sh (100%) rename mgeconvert/{onnx_converter => backend/ir_to_onnx}/onnx_converter.py (62%) create mode 100644 mgeconvert/backend/ir_to_onnx/onnx_op.py create mode 100644 mgeconvert/backend/ir_to_tflite/__init__.py create mode 100755 mgeconvert/backend/ir_to_tflite/init.sh rename mgeconvert/{tflite_converter => backend/ir_to_tflite}/pyflexbuffers/__init__.py (100%) rename mgeconvert/{tflite_converter => backend/ir_to_tflite}/pyflexbuffers/fbconverter.cc (100%) rename mgeconvert/{tflite_converter => backend/ir_to_tflite}/tflite_converter.py (66%) rename mgeconvert/{tflite_converter => backend/ir_to_tflite}/tflite_op.py (54%) delete mode 100644 mgeconvert/caffe_converter/caffe_converter.py delete mode 100644 mgeconvert/cambricon_converter/.gitignore delete mode 100644 mgeconvert/cambricon_converter/CMakeLists.txt delete mode 100644 mgeconvert/cambricon_converter/README.md delete mode 100644 mgeconvert/cambricon_converter/cambricon_converter.py delete mode 100644 mgeconvert/cambricon_converter/cambricon_op.py delete mode 100644 mgeconvert/cambricon_converter/cmake/cndev.cmake delete mode 100644 mgeconvert/cambricon_converter/cmake/cnml.cmake delete mode 100644 mgeconvert/cambricon_converter/cmake/cnrt.cmake delete mode 100755 mgeconvert/cambricon_converter/init.sh delete mode 100644 mgeconvert/cambricon_converter/lib/model.py delete mode 100644 mgeconvert/cambricon_converter/lib/operators.py delete mode 100644 mgeconvert/cambricon_converter/lib/tensor.py delete mode 100644 mgeconvert/cambricon_converter/swig/cambricon.i rename mgeconvert/{utils => converter_ir}/__init__.py (93%) create mode 100644 mgeconvert/converter_ir/ir_graph.py create mode 100644 mgeconvert/converter_ir/ir_op.py create mode 100644 mgeconvert/converter_ir/ir_quantizer.py create mode 100644 mgeconvert/converter_ir/ir_tensor.py create mode 100644 mgeconvert/converter_ir/ir_transform.py rename mgeconvert/{caffe_converter => converters}/__init__.py (83%) create mode 100644 mgeconvert/converters/mge_to_caffe.py create mode 100644 mgeconvert/converters/mge_to_onnx.py create mode 100644 mgeconvert/converters/mge_to_tflite.py create mode 100644 mgeconvert/converters/tm_to_caffe.py create mode 100644 mgeconvert/converters/tm_to_onnx.py create mode 100644 mgeconvert/converters/tm_to_tflite.py rename mgeconvert/{cambricon_converter => frontend}/__init__.py (81%) create mode 100644 mgeconvert/frontend/mge_to_ir/__init__.py create mode 100644 mgeconvert/frontend/mge_to_ir/mge_frontend.py rename mgeconvert/{mge_context => frontend/mge_to_ir}/mge_utils.py (99%) rename mgeconvert/{cambricon_converter/lib/cnlib => frontend/mge_to_ir/op_generators}/__init__.py (76%) create mode 100644 mgeconvert/frontend/mge_to_ir/op_generators/base.py create mode 100644 mgeconvert/frontend/mge_to_ir/op_generators/convolution.py create mode 100644 mgeconvert/frontend/mge_to_ir/op_generators/elemwise.py create mode 100644 mgeconvert/frontend/mge_to_ir/op_generators/tensor.py create mode 100644 mgeconvert/frontend/mge_to_ir/symbolvar_resolver.py create mode 100644 mgeconvert/frontend/tm_to_ir/__init__.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/__init__.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/base.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/batchnorm.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/broadcast.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/concat.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/constant.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/conv2d.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/deconv.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/dropout.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/elemwise.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/flatten.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/getvarshape.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/matmul.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/pooling.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/reduce.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/reshape.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/resize.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/softmax.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/squeeze.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/subtensor.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/transpose.py create mode 100644 mgeconvert/frontend/tm_to_ir/op_generators/typecvt.py create mode 100644 mgeconvert/frontend/tm_to_ir/pattern_utils.py create mode 100644 mgeconvert/frontend/tm_to_ir/qat_pattern.py create mode 100644 mgeconvert/frontend/tm_to_ir/tm_frontend.py create mode 100644 mgeconvert/frontend/tm_to_ir/tm_tensor_resolver.py create mode 100644 mgeconvert/frontend/tm_to_ir/tm_utils.py delete mode 100644 mgeconvert/mge_context/__init__.py delete mode 100644 mgeconvert/mge_context/mge_net.py delete mode 100644 mgeconvert/mge_context/mge_op.py delete mode 100644 mgeconvert/mge_context/mge_tensor.py delete mode 100644 mgeconvert/mge_context/mge_transform.py delete mode 100644 mgeconvert/onnx_converter/onnx_op.py delete mode 100644 mgeconvert/tflite_converter/__init__.py delete mode 100755 mgeconvert/tflite_converter/init.sh delete mode 100644 mgeconvert/utils/convert_caffe.py delete mode 100644 mgeconvert/utils/convert_cambricon.py delete mode 100644 mgeconvert/utils/convert_onnx.py delete mode 100644 mgeconvert/utils/convert_tflite.py delete mode 100644 test/convolution-backward-filter.mge create mode 100644 test/mge/__init__.py create mode 100644 test/mge/convolution-backward-filter.mge rename test/{ => mge}/test_caffe.py (97%) rename test/{ => mge}/test_onnx.py (91%) rename test/{ => mge}/test_tflite.py (76%) delete mode 100644 test/quantization_utils.py delete mode 100644 test/test_cambricon.py delete mode 100644 test/test_transform.py create mode 100644 test/traced_module/__init__.py create mode 100644 test/traced_module/test_caffe.py create mode 100644 test/traced_module/test_onnx.py create mode 100644 test/traced_module/test_qat_tflite.py create mode 100644 test/traced_module/test_tflite.py rename mgeconvert/cambricon_converter/lib/__init__.py => test/traced_module/tm_utils.py (59%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index eb1e17b..e8f71cf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,36 +28,6 @@ jobs: - name: lint run: ./ci/lint.sh - pytest-transform: - runs-on: ubuntu-latest - container: - image: enginesh233/mgeconvert_ci:v1.0 - defaults: - run: - shell: bash - steps: - - uses: actions/checkout@v2 - - name: test transform - run: ./ci/pytest_transform.sh - - pytest-cambricon: - runs-on: ubuntu-latest - container: - image: enginesh233/mgeconvert_ci:v1.0 - defaults: - run: - shell: bash - steps: - - uses: actions/checkout@v2 - - name: build cambricon - run: | - source ./ci/utils.sh - ./ci/build_cambricon.sh - - name: test cambricon - run: | - source ./ci/utils.sh - ./ci/pytest_cambricon.sh - pytest-caffe-and-onnx: runs-on: ubuntu-latest container: diff --git a/.gitignore b/.gitignore index 6fdf82d..a834910 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,8 @@ __pycache__ -mgeconvert/caffe_converter/caffe_pb -mgeconvert/tflite_converter/tflite/ +mgeconvert/backend/ir_to_caffe/caffe_pb +mgeconvert/backend/ir_to_tflite/tflite/ *.tflite *.onnx *.so .mypy_cache +*.tm \ No newline at end of file diff --git a/README.md b/README.md index e80a01d..57d1001 100644 --- a/README.md +++ b/README.md @@ -1,27 +1,42 @@ # MgeConvert -适用于 [MegEngine](https://github.com/MegEngine/MegEngine) 的各种转换器, 目前支持的框架有 [Caffe](https://github.com/BVLC/caffe)、[ONNX](https://github.com/onnx/onnx)、Cambricon 和 TFLite。 +适用于 [MegEngine](https://github.com/MegEngine/MegEngine) 的各种转换器, 目前支持的框架有 [Caffe](https://github.com/BVLC/caffe)、[ONNX](https://github.com/onnx/onnx) 和 TFLite。 -支持的模型包括 ResNet、ResNext、ShuffleNet 等，如果需要适配其他模型, 可能需要添加更多的算子支持。 +MgeConvert转换工具位于converters目录下，可直接调用其中的脚本将MegEngine导出的mge/TracedModule模型转换为第三方模型文件。 -MgeConvert 前端的部分位于 `mge_context` 目录下, 可以直接将 MegEngine dump 出来的计算图转为图结构, 方便后端语言生成。 +MgeConvert转换器的结构包含前端、中间表示（IR）、后端三个部分： +1. 前端的部分位于 `frontend` 目录下, 支持 mge 和 traced module 模型格式，可以将 MegEngine 序列化出来的计算图转为IR图结构 +2. IR部分位于 `converter_ir`目录下，包含图和 IR 算子定义、对计算图做变换的 transform rules 以及对量化模型处理的量化器 +3. 后端的部分位于 `backend` 目录下，包含caffe、onnx、tflite的转换器，可以将IR图结构转换为第三方框架的模型文件 +目前支持的模型包括 ResNet、ResNext、ShuffleNet 等，如果需要适配其他模型, 可能需要添加更多的算子支持。 ## 安装方式 -MgeConvert 基于 MegEngine 工作，因此确保您的电脑已经安装 MegEngine。 +MgeConvert 基于 MegEngine 工作，因此确保您的电脑已经安装 MegEngine(>=1.0)。 -以 caffe 为例，下面这条指令将安装 caffe 转换器并处理相关依赖。 +以 caffe 为例，下面这条指令将安装开发版本的 caffe 转换器并处理相关依赖。 ```bash python3 -m pip install git+https://github.com/MegEngine/mgeconvert.git --user --install-option="--targets=caffe" ``` -``--targets`` 的可选值有 ``caffe``、``onnx``、``cambricon``、``tflite`` 和 ``all``。 +``--targets`` 的可选值有 ``caffe``、``onnx``、``tflite`` 和 ``all``。 + +``tflite`` 转换器的schema默认使用r2.3版本，支持使用参数 ``tfversion`` 选择tflite schema的版本, 例如： + +```bash +--install-option="--targets=tflite --tfversion=r2.4" +``` ``all`` 代表安装全部转换器，不建议使用。可选值支持组合传入，比如 ``--targets=caffe,onnx`` 。 -> :warning: 由于 Cambricon SDK 未对外发布，寒武纪转换器将不会自动安装依赖，请事先查看依赖说明并配置环境。 +建议使用时指定版本号安装release版本的转换器，如安装0.4.2版本： + +```bash +python3 -m pip install git+https://github.com/MegEngine/mgeconvert.git@v0.4.2 --user --install-option="--targets=caffe" +``` +> :warning: 如果需要转换``TracedModule``模型，请安装v0.5.0以上版本 ## 使用方式 @@ -33,10 +48,50 @@ python3 -m pip install git+https://github.com/MegEngine/mgeconvert.git --user -- convert -h ``` -以 caffe 为例，查询转换参数： +以 mge模型转 caffe 为例，查询转换参数： ```bash -convert caffe -h +convert mge_to_caffe -h +``` + +### Feature 支持说明 + +- :white_check_mark: 已支持，并完成测试 +- :memo: 未支持，或尚未测试完全 +- :boom: 明确不支持 + +| TracedModule | tflite | caffe | onnx | +|---------------------|--------------------|--------------------|--------------------| +| QAT | :white_check_mark: | :memo: | :memo: | +| Quantized | :white_check_mark: | :boom: | :memo: | +| Float32 | :white_check_mark: | :white_check_mark: | :white_check_mark: | + +| Mgo | tflite | caffe | onnx | +|---------------------|--------------------|--------------------|--------------------| +| QAT | :boom: | :boom: | :boom: | +| Quantized | :memo: | :boom: | :memo: | +| Float32 | :white_check_mark: | :white_check_mark: | :white_check_mark: | + +### TFLite转换 +TFlite转换器支持 Float32 和量化的 TracedModule 转换。 +对于QAT模型，可以通过设置tracedmodule_to_tflite转换器中的 `require_quantize=True` 选择转换出tflite支持的量化数据类型（int8/uint8/int16/int32）量化后的Quantized 模型: + +```bash +convert tracedmodule_to_tflite -i tracedmodule.tm -o out.tflite --require_quantize +``` + +也可设置 `require_quantize=False` 选择转换出float32模型和量化参数文件。 + +```bash +convert tracedmodule_to_tflite -i tracedmodule.tm -o out.tflite --quantize_file_path "quant_params.json" +``` + +对于后者，还可以通过设置 `param_fake_quant` 参数来选择是否对参数进行假量化。 + +如果模型中没有QuantStub对输入数据进行量化处理，可以在转换时指定输入数据的量化类型、scale和zero_point量化参数 ： + +```bash +convert tracedmodule_to_tflite -i tracedmodule.tm -o out.tflite --input_data_type "quint8" --input_scales 0.125 --input_zero_points 128 --require_quantize ``` ## 依赖说明 @@ -49,11 +104,7 @@ convert caffe -h - Python packages: protobuf, onnx==1.7.0 -3. cambricon - - - Cambricon SDK: CNRT, CNML - -4. tflite +3. tflite - Python packages: pybind11==2.6.2, tensorflow==2.4.0 - third party: [flatbuffers](https://github.com/google/flatbuffers.git) @@ -61,36 +112,41 @@ convert caffe -h ## 算子支持列表 -| |Caffe|ONNX|Cambricon|TFLite| -|-- |-----|----|---------|------| -|abs| ✓ | ✓ | ✓ | ✓ | -|add| ✓ | ✓ | ✓ | ✓ | -|average pool2d| ✓ | ✓ | ✓ | ✓ | -|batchnorm| ✓ | ✓ | ✓ | × | -|broadcast| ✓ | ✓ | ✓ | × | -|ceil| × | ✓ | × | × | -|concat| ✓ | ✓ | ✓ | ✓ | -|conv2d| ✓ | ✓ | ✓ | ✓ | -|convtranspose2d| ✓ | ✓ | ✓ | ✓ | -|div(true_div)| ✓ | ✓ | ✓ | ✓ | -|exp| ✓ | ✓ | ✓ | ✓ | -|elemwise max| ✓ | ✓ | ✓ | ✓ | -|floor| × | ✓ | ✓ | × | -|log| ✓ | ✓ | ✓ | × | -|matrix mul| ✓ | ✓ | ✓ | ✓ | -|max pool2d| ✓ | ✓ | ✓ | ✓ | -|mul| ✓ | ✓ | ✓ | ✓ | -|pow| ✓ | ✓ | ✓ | ✓ | -|reduce max| ✓ | ✓ | ✓ | ✓ | -|reduce sum| ✓ | ✓ | ✓ | ✓ | -|relu| ✓ | ✓ | ✓ | ✓ | -|reshape| ✓ | ✓ | ✓ | ✓ | -|sigmoid| ✓ | ✓ | ✓ | × | -|softmax| ✓ | ✓ | ✓ | ✓ | -|leaky_relu| ✓ | × | × | ✓ | -|sub| ✓ | ✓ | ✓ | ✓ | -|slice(subtensor)| ✓ | ✓ | ✓ | ✓ | -|squeeze(axis_add_remove)| ✓ | ✓ | ✓ | ✓ | -|tanh| ✓ | ✓ | ✓ | ✓ | -|typecvt| ✓ | ✓ | ✓ | ✓ | -|transpose(dimshuffle)| ✓ | ✓ | ✓ | ✓ | +| tracemodule:rocket:
mgo:fire: | TFLite | Caffe | ONNX | +|--------------------------|---------|---------|---------| +| abs | ✓
✓ | ✓
✓ | ✓
✓ | +| average pool2d | ✓
✓ | ✓
✓ | ✓
✓ | +| batchnorm | ×
× | ✓
✓ | ✓
✓ | +| broadcast | ×
× | ✓
✓ | ✓
✓ | +| ceil | ×
× | ×
× | ✓
✓ | +| concat | ✓
✓ | ✓
✓ | ✓
✓ | +| conv2d | ✓
✓ | ✓
✓ | ✓
✓ | +| convtranspose2d | ✓
✓ | ✓
✓ | ✓
✓ | +| div(true_div) | ✓
✓ | ✓
✓ | ✓
✓ | +| exp | ✓
✓ | ✓
✓ | ✓
✓ | +| elemwise max | ✓
✓ | ✓
✓ | ✓
✓ | +| floor | ×
× | ×
× | ✓
✓ | +| log | ✓
✓ | ✓
✓ | ✓
✓ | +| matrix mul | ✓
✓ | ✓
✓ | ✓
✓ | +| max pool2d | ✓
✓ | ✓
✓ | ✓
✓ | +| mul | ✓
✓ | ✓
✓ | ✓
✓ | +| pow | ✓
✓ | ✓
✓ | ✓
✓ | +| reduce max | ✓
✓ | ✓
✓ | ✓
✓ | +| reduce min | ✓
✓ | ✓
✓ | ✓
✓ | +| reduce mean | ✓
✓ | ✓
✓ | ✓
✓ | +| reduce sum | ✓
✓ | ✓
✓ | ✓
✓ | +| relu | ✓
✓ | ✓
✓ | ✓
✓ | +| relu6 | ✓
✓ | ×
× | ✓
✓ | +| reshape | ✓
✓ | ✓
✓ | ✓
✓ | +| resize | ✓
✓ | ✓
✓ | ✓
✓ | +| sigmoid(logistic) | ✓
✓ | ✓
✓ | ✓
✓ | +| softmax | ✓
✓ | ✓
✓ | ✓
✓ | +| leaky_relu | ✓
✓ | ✓
✓ | ✓
✓ | +| sub | ✓
✓ | ✓
✓ | ✓
✓ | +| slice(subtensor) | ✓
✓ | ✓
✓ | ✓
✓ | +| squeeze(axis_add_remove) | ✓
✓ | ✓
✓ | ✓
✓ | +| tanh | ✓
✓ | ✓
✓ | ✓
✓ | +| typecvt | ✓
✓ | ✓
✓ | ✓
✓ | +| transpose(dimshuffle) | ✓
✓ | ✓
✓ | ✓
✓ | +| AdaptiveAvgPool2d | ×
× | ×
× | ✓
✓ | +| flatten | ×
× | ×
× | ✓
✓ | diff --git a/bin/convert b/bin/convert index 807b410..32e9795 100755 --- a/bin/convert +++ b/bin/convert @@ -1,6 +1,7 @@ #!/usr/bin/python3 import argparse +from typing import List, Union import mgeconvert @@ -8,141 +9,167 @@ import mgeconvert def get_targets(module): targets = [] for attr in dir(module): - if attr.startswith("convert_to"): - targets.append(attr[11:]) + if attr.startswith("mge_to") or attr.startswith("tracedmodule_to"): + targets.append(attr) return targets def init(subparsers): targets = get_targets(mgeconvert) - if "caffe" in targets: - def to_caffe(args): - outspec = None - if args.end_point is not None: - outspec = args.end_point.split(";") - mgeconvert.convert_to_caffe( - args.input, prototxt=args.prototxt, caffemodel=args.caffemodel, outspec=outspec - ) - def caffe_parser(subparsers): - p = subparsers.add_parser("caffe",) - p.set_defaults(func=to_caffe) - p.add_argument( - "-i", "--input", required=True, type=str, help="Input megengine dump model file" - ) - p.add_argument( - "-c", "--prototxt", required=True, type=str, help="Output caffe .prototxt file" - ) - p.add_argument( - "-b", - "--caffemodel", - required=True, - type=str, - help="Output caffe .caffemodel file", - ) + for target in targets: + if "caffe" in target: + def to_caffe(args): + outspec = None + if args.end_point is not None: + outspec = args.end_point.split(";") + converter_map = { + "tracedmodule_to_caffe": mgeconvert.tracedmodule_to_caffe, + "mge_to_caffe": mgeconvert.mge_to_caffe, + } + converter_map[target]( + args.input, prototxt=args.prototxt, caffemodel=args.caffemodel, outspec=outspec + ) + def caffe_parser(subparsers): + p = subparsers.add_parser(target) + p.set_defaults(func=to_caffe) + p.add_argument( + "-i", "--input", required=True, type=str, help="Input megengine dump model file" + ) + p.add_argument( + "-c", "--prototxt", required=True, type=str, help="Output caffe .prototxt file" + ) + p.add_argument( + "-b", + "--caffemodel", + required=True, + type=str, + help="Output caffe .caffemodel file", + ) - p.add_argument( - "--end_point", - default=None, - type=str, - help="end_point is used to specify which part of the mge model should be converted", - ) - caffe_parser(subparsers) - if "cambricon" in targets: - def to_cambricon(args): - mgeconvert.convert_to_cambricon( - args.input, - args.output, - args.batch_size, - args.core_number, - args.data_type, - args.use_nhwc, - ) - def cambricon_parser(subparsers): - p = subparsers.add_parser("cambricon") - p.set_defaults(func=to_cambricon) - p.add_argument( - "-i", "--input", required=True, type=str, help="megengine dumped model file" - ) - p.add_argument( - "-o", "--output", required=True, type=str, help="converted Cambricon model file" - ) - p.add_argument( - "-b", "--batch-size", default=4, type=int, help="best practice: 4" - ) - p.add_argument("-c", "--core-number", default=1, type=int, help="c <= 16") - p.add_argument( - "-t", "--data-type", default="float32", type=str, help="float32, float16" - ) - p.add_argument("--use-nhwc", action="store_true", help="default nchw") - cambricon_parser(subparsers) - if "onnx" in targets: - def to_onnx(args): - outspec = None - if args.end_point is not None: - outspec = args.end_point.split(";") + p.add_argument( + "--end_point", + default=None, + type=str, + help="end_point is used to specify which part of the mge model should be converted", + ) + caffe_parser(subparsers) + if "onnx" in target: + def to_onnx(args): + outspec = None + if args.end_point is not None: + outspec = args.end_point.split(";") + converter_map = { + "tracedmodule_to_onnx": mgeconvert.tracedmodule_to_onnx, + "mge_to_onnx": mgeconvert.mge_to_onnx, + } + converter_map[target]( + args.input, + args.output, + graph_name=args.graph, + opset=args.opset, + outspec=outspec, + ) + def onnx_parser(subparsers): + p = subparsers.add_parser("onnx") + p.set_defaults(func=to_onnx) + p.add_argument( + "-i", "--input", required=True, type=str, help="Input megengine dump model file" + ) + p.add_argument( + "-o", "--output", required=True, type=str, help="Output onnx .onnx file" + ) + p.add_argument("--opset", default=8, type=int, help="Onnx opset version") + p.add_argument("--graph_name", default="graph", type=str, help="Onnx graph name") + p.add_argument( + "--end_point", + default=None, + type=str, + help="end_point is used to specify which part of the mge model should be converted", + ) + onnx_parser(subparsers) + if "tflite" in target: + def to_tflite(args): + if "tracedmodule" in target: + mgeconvert.tracedmodule_to_tflite( + mge_fpath=args.input, + output=args.output, + input_data_type = args.input_data_type, + input_scales = args.input_scales, + input_zero_points = args.input_zero_points, + require_quantize=args.require_quantize, + param_fake_quant=args.param_fake_quant, + quantize_file_path=args.quantize_file_path, + graph_name=args.graph_name, + mtk=args.mtk, + ) + else: + mgeconvert.mge_to_tflite( + mge_fpath=args.input, + output=args.output, + graph_name=args.graph_name, + mtk=args.mtk, + ) + def tflite_parser(subparsers): + p = subparsers.add_parser(target) + p.set_defaults(func=to_tflite) + p.add_argument( + "-i", "--input", required=True, type=str, help="megengine dumped model file" + ) + p.add_argument( + "-o", "--output", required=True, type=str, help="converted TFLite model file" + ) + if "tracedmodule" in target: + p.add_argument( + "--input_data_type", + default=None, + type=str, + help="the dtype of input data used for quantize model input", + ) + p.add_argument( + "--input_scales", + default=None, + type=Union[float, List[float]], + help="the scale of input data used for quantize model input", + ) + p.add_argument( + "--input_zero_points", + default=None, + type=Union[int, List[int]], + help="the zero point of input data used for quantize model input", + ) + p.add_argument( + "--require_quantize", + action="store_false", + help="whether to do quantize if the model has quantization parameters", + ) + p.add_argument( + "--param_fake_quant", + action="store_false", + help="whether to do fake quantize for parameters if the model has quantization parameters", + ) + p.add_argument( + "--quantize_file_path", + default="quant_params.json", + type=str, + help="", + ) + p.add_argument( + "--graph-name", + default="graph0", + type=str, + help="default subgraph name in TFLite model", + ) + p.add_argument( + "--mtk", action="store_true", help="If target flatform is MTK(P70, P80)" + ) - mgeconvert.convert_to_onnx( - args.input, - args.output, - graph_name=args.graph, - opset=args.opset, - outspec=outspec, - ) - def onnx_parser(subparsers): - p = subparsers.add_parser("onnx") - p.set_defaults(func=to_onnx) - p.add_argument( - "-i", "--input", required=True, type=str, help="Input megengine dump model file" - ) - p.add_argument( - "-o", "--output", required=True, type=str, help="Output onnx .onnx file" - ) - p.add_argument("--opset", default=8, type=int, help="Onnx opset version") - p.add_argument("--graph", default="graph", type=str, help="Onnx graph name") - p.add_argument( - "--end_point", - default=None, - type=str, - help="end_point is used to specify which part of the mge model should be converted", - ) - onnx_parser(subparsers) - if "tflite" in targets: - def to_tflite(args): - mgeconvert.convert_to_tflite( - mge_fpath=args.input, - output=args.output, - graph_name=args.graph_name, - batch_size=args.batch_size, - mtk=args.mtk, - ) - def tflite_parser(subparsers): - p = subparsers.add_parser("tflite") - p.set_defaults(func=to_tflite) - p.add_argument( - "-i", "--input", required=True, type=str, help="megengine dumped model file" - ) - p.add_argument( - "-o", "--output", required=True, type=str, help="converted TFLite model file" - ) - p.add_argument( - "--graph-name", - default="graph0", - type=str, - help="default subgraph name in TFLite model", - ) - p.add_argument( - "-b", "--batch-size", default=1, type=int, help="default value: 1" - ) - p.add_argument( - "--mtk", action="store_true", help="If target flatform is MTK(P70, P80)" - ) - tflite_parser(subparsers) + tflite_parser(subparsers) def main(): targets = get_targets(mgeconvert) - msg = targets[0] if len(targets) == 1 else "{" + ",".join(targets) + "}" + msg = targets[0] if len(targets) == 1 else "{" + ", ".join(targets) + "}" parser = argparse.ArgumentParser( description="use \"convert %s -h\" for more details" % msg ) diff --git a/ci/build_cambricon.sh b/ci/build_cambricon.sh deleted file mode 100755 index f90950e..0000000 --- a/ci/build_cambricon.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash -e - -set -e - -wget -P mgeconvert/cambricon_converter/swig https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/numpy.i - -mkdir -p mgeconvert/cambricon_converter/build -cd mgeconvert/cambricon_converter/build -cmake .. -make -j4 -make develop -cd ../../.. diff --git a/ci/pytest_caffe_and_onnx.sh b/ci/pytest_caffe_and_onnx.sh index 4a62918..59673f1 100755 --- a/ci/pytest_caffe_and_onnx.sh +++ b/ci/pytest_caffe_and_onnx.sh @@ -8,34 +8,45 @@ python3 -m pip install --no-binary=protobuf protobuf==3.8.0 apt install -y protobuf-compiler -./mgeconvert/caffe_converter/init.sh +./mgeconvert/backend/ir_to_caffe/init.sh pip3 install scikit-image==0.17.2 +sudo -H python3 -m pip install -q megengine==1.6.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py +pytest test/traced_module/test_caffe.py +pytest test/traced_module/test_onnx.py +sudo -H python3 -m pip uninstall -y megengine + +sudo -H python3 -m pip install -q megengine==1.5.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py +sudo -H python3 -m pip uninstall -y megengine + +sudo -H python3 -m pip install -q megengine==1.4.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py +sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.3.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_caffe.py -pytest test/test_onnx.py +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.2.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_caffe.py -pytest test/test_onnx.py +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.1.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_caffe.py -pytest test/test_onnx.py +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.0.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_caffe.py -pytest test/test_onnx.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==0.6.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_caffe.py -pytest test/test_onnx.py +pytest test/mge/test_caffe.py +pytest test/mge/test_onnx.py sudo -H python3 -m pip uninstall -y megengine \ No newline at end of file diff --git a/ci/pytest_cambricon.sh b/ci/pytest_cambricon.sh deleted file mode 100755 index 88561e1..0000000 --- a/ci/pytest_cambricon.sh +++ /dev/null @@ -1,23 +0,0 @@ -#!/bin/bash -e - -set -e - -sudo -H python3 -m pip install -q megengine==1.3.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_cambricon.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.2.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_cambricon.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.1.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_cambricon.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.0.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_cambricon.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==0.6.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_cambricon.py -sudo -H python3 -m pip uninstall -y megengine \ No newline at end of file diff --git a/ci/pytest_tflite.sh b/ci/pytest_tflite.sh index 13d2956..b5a55fd 100755 --- a/ci/pytest_tflite.sh +++ b/ci/pytest_tflite.sh @@ -2,30 +2,40 @@ set -e -./mgeconvert/tflite_converter/init.sh +./mgeconvert/backend/ir_to_tflite/init.sh # try to find libflatbuffers.so export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib sudo python3 -m pip uninstall flatbuffers -y -sudo python3 -m pip install tensorflow==2.4.0 +sudo python3 -m pip install tensorflow==2.5.0 + +sudo -H python3 -m pip install -q megengine==1.6.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_tflite.py +pytest test/traced_module/test_tflite.py +pytest test/traced_module/test_qat_tflite.py +sudo -H python3 -m pip uninstall -y megengine + +sudo -H python3 -m pip install -q megengine==1.5.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_tflite.py +sudo -H python3 -m pip uninstall -y megengine + +sudo -H python3 -m pip install -q megengine==1.4.0 -f https://megengine.org.cn/whl/mge.html +pytest test/mge/test_tflite.py +sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.3.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_tflite.py +pytest test/mge/test_tflite.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.2.0 -f https://megengine.org.cn/whl/mge.html -pytest -v test/test_tflite.py +pytest -v test/mge/test_tflite.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.1.0 -f https://megengine.org.cn/whl/mge.html -pytest -v test/test_tflite.py +pytest -v test/mge/test_tflite.py sudo -H python3 -m pip uninstall -y megengine sudo -H python3 -m pip install -q megengine==1.0.0 -f https://megengine.org.cn/whl/mge.html -pytest -v test/test_tflite.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==0.6.0 -f https://megengine.org.cn/whl/mge.html -pytest -v test/test_tflite.py -sudo -H python3 -m pip uninstall -y megengine +pytest -v test/mge/test_tflite.py +sudo -H python3 -m pip uninstall -y megengine \ No newline at end of file diff --git a/ci/pytest_transform.sh b/ci/pytest_transform.sh deleted file mode 100755 index a9afc41..0000000 --- a/ci/pytest_transform.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/bin/bash -e - -set -e - -python3 -m pip install -q -r ci/requires-test.txt - -sudo -H python3 -m pip install -q megengine==1.3.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_transform.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.2.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_transform.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.1.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_transform.py -sudo -H python3 -m pip uninstall -y megengine - -sudo -H python3 -m pip install -q megengine==1.0.0 -f https://megengine.org.cn/whl/mge.html -pytest test/test_transform.py -sudo -H python3 -m pip uninstall -y megengine diff --git a/ci/utils.sh b/ci/utils.sh deleted file mode 100644 index a917ac8..0000000 --- a/ci/utils.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/bin/bash -e - -export NEUWARE_HOME=/usr/local/neuware -export LD_LIBRARY_PATH=/usr/local/neuware/lib64 -export LIBRARY_PATH=/usr/local/neuware/lib64 diff --git a/mgeconvert/__init__.py b/mgeconvert/__init__.py index 1207b5d..93515f3 100644 --- a/mgeconvert/__init__.py +++ b/mgeconvert/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. diff --git a/mgeconvert/backend/__init__.py b/mgeconvert/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/mgeconvert/onnx_converter/__init__.py b/mgeconvert/backend/ir_to_caffe/__init__.py similarity index 83% rename from mgeconvert/onnx_converter/__init__.py rename to mgeconvert/backend/ir_to_caffe/__init__.py index d8edbf5..425e87b 100644 --- a/mgeconvert/onnx_converter/__init__.py +++ b/mgeconvert/backend/ir_to_caffe/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,4 +5,4 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .onnx_converter import convert_to_onnx +from .caffe_converter import CaffeConverter diff --git a/mgeconvert/backend/ir_to_caffe/caffe_converter.py b/mgeconvert/backend/ir_to_caffe/caffe_converter.py new file mode 100644 index 0000000..8c70962 --- /dev/null +++ b/mgeconvert/backend/ir_to_caffe/caffe_converter.py @@ -0,0 +1,117 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error +from google.protobuf import text_format # type: ignore +from tqdm import tqdm + +from ...converter_ir.ir_graph import IRGraph +from .caffe_op import MGE2CAFFE, _add_input_layer +from .caffe_pb import caffe_pb2 as cp + + +class CaffeConverter: + def __init__(self, net, use_empty_blobs=False): + assert isinstance(net, IRGraph), "net must be instance of IRGraph" + self.net = net + self.tensor2blob_map = {} + self.layers = [] + self._names = set() + self._count = 0 + self.use_empty_blobs = use_empty_blobs + + def dump(self, proto_file, caffe_file=None): + CaffeNet = cp.NetParameter(layer=self.layers) + if caffe_file is not None: + with open(caffe_file, "wb") as f: + f.write(CaffeNet.SerializeToString()) + + for layer in CaffeNet.layer: + layer.ClearField("blobs") + + with open(proto_file, "w", encoding=None) as f: + f.write(text_format.MessageToString(CaffeNet)) + + @property + def gen_name(self): + self._count = self._count + 1 + while "_caffe_{0}".format(self._count) in self._names: + self._count = self._count + 1 + return "_caffe_{0}".format(self._count) + + def get_blob_name(self, tensor): + if tensor not in self.tensor2blob_map: + raise KeyError("can not find tensor {}".format(tensor)) + return self.tensor2blob_map[tensor] + + def set_blob_name(self, tensor, name=None): + assert tensor not in self.tensor2blob_map, "{} already be set".format(tensor) + if name is not None: + assert isinstance(name, str) + self.tensor2blob_map[tensor] = name + else: + self.tensor2blob_map[tensor] = self.gen_name + self._names.add(self.tensor2blob_map[tensor]) + return self.tensor2blob_map[tensor] + + def reset_blob_name(self, tensor, name=None): + assert tensor in self.tensor2blob_map, "{} should be set".format(tensor) + if name is not None: + assert isinstance(name, str) + self.tensor2blob_map[tensor] = name + else: + self.tensor2blob_map[tensor] = self.gen_name + self._names.add(self.tensor2blob_map[tensor]) + return self.tensor2blob_map[tensor] + + def gen_blob_proto(self, data): + if self.use_empty_blobs: + return cp.BlobProto() + if isinstance(data, (int, float)): + return cp.BlobProto(data=[data]) + else: + return cp.BlobProto( + data=data.reshape(-1), shape=cp.BlobShape(dim=data.shape) + ) + + def add_layer(self, layer): + if isinstance(layer, list): + for x in layer: + self.layers.append(x) + else: + self.layers.append(layer) + + def convert(self): + unsupported_oprs = [] + for opr in self.net.all_oprs: + if not isinstance(opr, tuple(MGE2CAFFE.keys())): + unsupported_oprs.append(opr) + continue + unsupported_oprs = set(map(type, unsupported_oprs)) + assert not unsupported_oprs, "Operators {} are not supported yet".format( + unsupported_oprs + ) + + def need_convert(mge_opr): + is_const = [data.np_data is not None for data in mge_opr.inp_tensors] + return not all(is_const) and len(mge_opr.inp_tensors) > 0 + + all_oprs = list(self.net.all_oprs) + + # add net input + for net_inp in self.net.graph_inputs: + _add_input_layer(net_inp, self) + + for index in range(len(all_oprs) - 1, -1, -1): + if all_oprs[index].skip: + del all_oprs[index] + + for opr in tqdm(all_oprs): + if not need_convert(opr): + continue + MGE2CAFFE[type(opr)](opr, self) diff --git a/mgeconvert/caffe_converter/caffe_op.py b/mgeconvert/backend/ir_to_caffe/caffe_op.py similarity index 51% rename from mgeconvert/caffe_converter/caffe_op.py rename to mgeconvert/backend/ir_to_caffe/caffe_op.py index 25ac5c4..11705ca 100644 --- a/mgeconvert/caffe_converter/caffe_op.py +++ b/mgeconvert/backend/ir_to_caffe/caffe_op.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,48 +5,74 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# pylint: disable=import-error import collections +import os from math import ceil +from typing import Sequence import numpy as np +from megengine import get_logger -from ..mge_context import ( +from ...converter_ir.ir_op import ( + AbsOpr, + AdaptiveAvgPool2dOpr, + AddOpr, + AvgPool2dOpr, AxisAddRemoveOpr, - BatchNormForwardOpr, + BatchNormalizationOpr, BroadcastOpr, ConcatOpr, - ConvForwardBiasOpr, - ConvolutionBackwardDataBiasOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - DimshuffleOpr, - ElemwiseOpr, - FullyConnectedOpr, + ConstantOpr, + Conv2dOpr, + Deconv2dOpr, + ExpOpr, + FlattenOpr, + GetSubTensorOpr, GetVarShapeOpr, - Host2DeviceCopyOpr, + HardSigmoidOpr, + HardSwishOpr, IdentityOpr, LeakyReluOpr, - MarkNoBroadcastElemwiseOpr, - MatrixMulOpr, + LinearOpr, + LinspaceOpr, + LogOpr, + MatMulOpr, + MaxOpr, + MaxPool2dOpr, + MinOpr, + MulOpr, MultipleDeviceTensorHolderOpr, - PoolingForwardOpr, + OpBase, + PowOpr, ReduceOpr, + Relu6Opr, + ReluOpr, ReshapeOpr, SharedDeviceTensorOpr, - SubtensorOpr, - Tensor, + SigmoidOpr, + SiLUOpr, + SoftmaxOpr, + SqueezeOpr, + SubOpr, + TanHOpr, + TransposeOpr, + TrueDivOpr, TypeCvtOpr, VolatileSharedDeviceTensorOpr, - get_logger, - get_symvar_value, ) -from ..mge_context.mge_utils import isconst -from .caffe_pb import caffe_pb2 as cp # pylint: disable=import-error +from ...converter_ir.ir_tensor import IRTensor +from ...frontend.mge_to_ir.mge_utils import get_symvar_value +from .caffe_pb import caffe_pb2 as cp logger = get_logger(__name__) MGE2CAFFE = {} +def isconst(x): + return x.np_data is not None + + def _register_op(*oprs): def callback(impl): for opr in oprs: @@ -63,15 +88,28 @@ def silence_blob(blob): ) +def _add_input_layer(tensor, context): + param = cp.InputParameter(shape=[cp.BlobShape(dim=tensor.shape)]) + context.add_layer( + cp.LayerParameter( + bottom=[], + top=[context.set_blob_name(tensor, tensor.name)], + name="data_input", + type="Input", + input_param=param, + ) + ) + + def _gen_layer(opr, etype, context, single_input=True, **kwargs): bottom = ( - [context.get_blob_name(opr.inp_vars[0])] + [context.get_blob_name(opr.inp_tensors[0])] if single_input - else list(map(context.get_blob_name, opr.inp_vars)) + else list(map(context.get_blob_name, opr.inp_tensors)) ) - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] return cp.LayerParameter( - bottom=bottom, top=top, name=opr.name, type=etype, **kwargs + bottom=bottom, top=top, name=opr.out_tensors[0].name, type=etype, **kwargs ) @@ -179,13 +217,32 @@ def _broadcast_for_eltwiseopr(oprname, InputA, InputB, context): return (topA, topB, shape) -@_register_op(DimshuffleOpr) -def _dimshufulle(opr, context): +@_register_op( + MultipleDeviceTensorHolderOpr, + SharedDeviceTensorOpr, + VolatileSharedDeviceTensorOpr, + LinspaceOpr, +) +def _ignore(*_): + pass + + +@_register_op(GetVarShapeOpr) +def shapeof(opr): + out_shape = opr.out_tensors[0] + if hasattr(out_shape, "_var"): + out_shape.np_data = get_symvar_value(out_shape._var) + else: + out_shape.np_data = np.array(opr.inp_tensors[0].shape, dtype=np.int64) + + +@_register_op(TransposeOpr) +def _dimshfulle(opr, context): def swap_two_dimension(x, y, f, top, context): - prefix = "{}_swap_{}_{}".format(opr.name, x, y) - if opr.inp_vars[0].shape[f[x]] > 1: + prefix = "{}_swap_{}_{}".format(opr.out_tensors[0].name, x, y) + if opr.inp_tensors[0].shape[f[x]] > 1: bottom = top - top = [context.gen_name for _ in range(opr.inp_vars[0].shape[f[x]])] + top = [context.gen_name for _ in range(opr.inp_tensors[0].shape[f[x]])] context.add_layer( cp.LayerParameter( name="{}_slice_x".format(prefix), @@ -193,15 +250,16 @@ def swap_two_dimension(x, y, f, top, context): bottom=bottom, top=top, slice_param=cp.SliceParameter( - axis=x, slice_point=list(range(1, opr.inp_vars[0].shape[f[x]])) + axis=x, + slice_point=list(range(1, opr.inp_tensors[0].shape[f[x]])), ), ) ) - if opr.inp_vars[0].shape[f[y]] > 1: + if opr.inp_tensors[0].shape[f[y]] > 1: bottom = top top = [ - [context.gen_name for _ in range(opr.inp_vars[0].shape[f[y]])] - for _ in range(opr.inp_vars[0].shape[f[x]]) + [context.gen_name for _ in range(opr.inp_tensors[0].shape[f[y]])] + for _ in range(opr.inp_tensors[0].shape[f[x]]) ] context.add_layer( list( @@ -213,15 +271,17 @@ def swap_two_dimension(x, y, f, top, context): top=top[z], slice_param=cp.SliceParameter( axis=y, - slice_point=list(range(1, opr.inp_vars[0].shape[f[y]])), + slice_point=list( + range(1, opr.inp_tensors[0].shape[f[y]]) + ), ), ), - range(opr.inp_vars[0].shape[f[x]]), + range(opr.inp_tensors[0].shape[f[x]]), ) ) ) bottom = top - top = [context.gen_name for _ in range(opr.inp_vars[0].shape[f[x]])] + top = [context.gen_name for _ in range(opr.inp_tensors[0].shape[f[x]])] context.add_layer( list( map( @@ -232,11 +292,11 @@ def swap_two_dimension(x, y, f, top, context): top=[top[z]], concat_param=cp.ConcatParameter(axis=x), ), - range(opr.inp_vars[0].shape[f[x]]), + range(opr.inp_tensors[0].shape[f[x]]), ) ) ) - if opr.inp_vars[0].shape[f[x]] > 1: + if opr.inp_tensors[0].shape[f[x]] > 1: bottom = top top = [context.gen_name] context.add_layer( @@ -254,7 +314,7 @@ def swap_two_dimension(x, y, f, top, context): logger.warning("Add 'slice/concat layers' in operator: Dimshuffle") l = list(filter(lambda x: x != "x", list(opr.pattern))) - assert len(l) == opr.inp_vars[0].ndim + assert len(l) == opr.inp_tensors[0].ndim nl = [] for i, _ in enumerate(l): while l[i] != i: @@ -265,29 +325,29 @@ def swap_two_dimension(x, y, f, top, context): f = list(range(len(l))) - top = [context.get_blob_name(opr.inp_vars[0])] + top = [context.get_blob_name(opr.inp_tensors[0])] for i in nl: top = swap_two_dimension(i[0], i[1], f, top, context) - if len(opr.pattern) != opr.inp_vars[0].ndim: + if len(opr.pattern) != opr.inp_tensors[0].ndim: bottom = top top = [context.gen_name] context.add_layer( cp.LayerParameter( - name="{}_reshape".format(opr.name), + name="{}_reshape".format(opr.out_tensors[0].name), type="Reshape", bottom=bottom, top=top, reshape_param=cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.shapes)), ) ) - context.set_blob_name(opr.out_vars[0], list(top)[0]) + context.set_blob_name(opr.out_tensors[0], list(top)[0]) -@_register_op(SubtensorOpr) -def _subtensor(opr: SubtensorOpr, context): +@_register_op(GetSubTensorOpr) +def _subtensor(opr, context): logger.warning("Add 'slice/concat layers' in operator: Subtensor") - top = [context.get_blob_name(opr.inp_vars[0])] + top = [context.get_blob_name(opr.inp_tensors[0])] def axis_suffix_gen(): yield "" @@ -320,14 +380,14 @@ def top_name(k, concat_list, name): return name return name + context.gen_name - for i in range(0, len(opr.begin_param)): - begin = opr.begin_param[i] - end = opr.end_param[i] - step = opr.step_param[i] + for i in range(0, len(opr.begin_params)): + begin = opr.begin_params[i] + end = opr.end_params[i] + step = opr.step_params[i] axis = opr.axis[i] sl = slice(begin, end, step) - n = opr.inp_vars[0].shape[axis] + n = opr.inp_tensors[0].shape[axis] slice_list = get_slice_list(n, sl, step < 0) if slice_list == []: @@ -337,7 +397,7 @@ def top_name(k, concat_list, name): continue bottom = top - name = opr.name + next(axis_suffixes) + name = opr.out_tensors[0].name + next(axis_suffixes) spldict = collections.OrderedDict( [(k, top_name(k, concat_list, name)) for k in [0] + slice_list] @@ -377,13 +437,13 @@ def top_name(k, concat_list, name): ) ) - context.set_blob_name(opr.out_vars[0], list(top)[0]) + context.set_blob_name(opr.out_tensors[0], list(top)[0]) if len(opr.squeeze_axis): - name = opr.name + "_reshape" - param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_vars[0].shape)) + name = opr.out_tensors[0].name + "_reshape" + param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_tensors[0].shape)) bottom = top - top = [context.reset_blob_name(opr.out_vars[0], name)] + top = [context.reset_blob_name(opr.out_tensors[0], name)] context.add_layer( cp.LayerParameter( name=name, type="Reshape", bottom=bottom, top=top, reshape_param=param @@ -392,27 +452,6 @@ def top_name(k, concat_list, name): logger.warning("Add 'reshape layers' in operator: Subtensor") -@_register_op( - MultipleDeviceTensorHolderOpr, SharedDeviceTensorOpr, VolatileSharedDeviceTensorOpr -) -def _(*_): - pass - - -@_register_op(GetVarShapeOpr) -def shapeof(opr, _): - out_shape = opr.out_vars[0] - out_shape.np_data = get_symvar_value(out_shape._var) - - -@_register_op(Host2DeviceCopyOpr) -def _data_provider(opr, context): - param = cp.InputParameter(shape=[cp.BlobShape(dim=opr.shape)]) - context.add_layer( - _gen_layer(opr, "Input", context, single_input=False, input_param=param) - ) - - def bias_add(input, output, bias, name, context): param = cp.BiasParameter(axis=1, num_axes=1) blobs = [context.gen_blob_proto(bias)] @@ -431,7 +470,6 @@ def bias_add(input, output, bias, name, context): def _arith_with_const_tensor(input, const, order, opr, context): - atype = opr.mode topB = const.np_data if input.ndim >= 2 and ( topB.squeeze().shape == (input.shape[1],) or topB.squeeze().shape == (1,) @@ -441,27 +479,29 @@ def _arith_with_const_tensor(input, const, order, opr, context): shape = topB.shape layer_param = cp.ScaleParameter() else: - topA, topB, shape = _broadcast_for_eltwiseopr(opr.name, input, const, context) + topA, topB, shape = _broadcast_for_eltwiseopr( + opr.out_tensors[0].name, input, const, context + ) layer_param = cp.ScaleParameter(axis=len(shape) - topB.ndim, num_axes=topB.ndim) - if atype in {"ADD", "SUB"}: + if isinstance(opr, (AddOpr, SubOpr)): layer_param.bias_term = True param_b = topB param_k = np.ones(shape=param_b.shape) - if atype == "SUB": + if isinstance(opr, SubOpr): if order == 0: param_b = -param_b # pylint: disable=invalid-unary-operand-type else: - param_k = -param_k + param_k = -param_k # pylint: disable=invalid-unary-operand-type blobs = [context.gen_blob_proto(param_k), context.gen_blob_proto(param_b)] else: param_k = topB - if atype == "TRUE_DIV": + if isinstance(opr, TrueDivOpr): if order == 0: param_k = 1.0 / param_k else: bottom = topA - name = opr.name + context.gen_name + name = opr.out_tensors[0].name + context.gen_name topA = [name] context.add_layer( cp.LayerParameter( @@ -474,10 +514,10 @@ def _arith_with_const_tensor(input, const, order, opr, context): ) blobs = [context.gen_blob_proto(param_k)] bottom = topA - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( - name=opr.name, + name=opr.out_tensors[0].name, type="Scale", bottom=bottom, top=top, @@ -487,42 +527,42 @@ def _arith_with_const_tensor(input, const, order, opr, context): ) -def _arith(opr, mode, context): - atype = mode - if isconst(opr.inp_vars[0]) and isconst(opr.inp_vars[1]): +def _arith(opr, context): + assert isinstance(opr, (AddOpr, SubOpr, TrueDivOpr, MulOpr)) + + if isconst(opr.inp_tensors[0]) and isconst(opr.inp_tensors[1]): return - elif isconst(opr.inp_vars[0]) or isconst(opr.inp_vars[1]): - if isconst(opr.inp_vars[0]): - inpA = opr.inp_vars[1] - const = opr.inp_vars[0] + elif isconst(opr.inp_tensors[0]) or isconst(opr.inp_tensors[1]): + if isconst(opr.inp_tensors[0]): + inpA = opr.inp_tensors[1] + const = opr.inp_tensors[0] order = 1 else: - inpA = opr.inp_vars[0] - const = opr.inp_vars[1] + inpA = opr.inp_tensors[0] + const = opr.inp_tensors[1] order = 0 use_bias_layer = False bias = 0 - if atype == "ADD": + if isinstance(opr, AddOpr): bias = const.np_data.squeeze() if bias.ndim == 1 and inpA.ndim > 1 and bias.shape[0] == inpA.shape[1]: use_bias_layer = True if use_bias_layer: - bias_add(inpA, opr.out_vars[0], bias, opr.out_vars[0].name, context) + bias_add(inpA, opr.out_tensors[0], bias, opr.out_tensors[0].name, context) else: _arith_with_const_tensor(inpA, const, order, opr, context) - else: topA, topB, _ = _broadcast_for_eltwiseopr( - opr.name, opr.inp_vars[0], opr.inp_vars[1], context + opr.out_tensors[0].name, opr.inp_tensors[0], opr.inp_tensors[1], context ) - if atype in {"ADD", "SUB"}: + if isinstance(opr, (AddOpr, SubOpr)): param = cp.EltwiseParameter(operation="SUM") - if atype == "SUB": + if isinstance(opr, SubOpr): param.coeff.extend([1, -1]) else: - if atype == "TRUE_DIV": + if isinstance(opr, TrueDivOpr): bottom = topB - name = opr.name + context.gen_name + name = opr.out_tensors[0].name + context.gen_name topB = [name] context.add_layer( cp.LayerParameter( @@ -535,10 +575,10 @@ def _arith(opr, mode, context): ) param = cp.EltwiseParameter(operation="PROD") bottom = topA + topB - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( - name=opr.name, + name=opr.out_tensors[0].name, type="Eltwise", bottom=bottom, top=top, @@ -547,94 +587,51 @@ def _arith(opr, mode, context): ) -@_register_op(ElemwiseOpr, MarkNoBroadcastElemwiseOpr, IdentityOpr) -def _eltwise(opr: ElemwiseOpr, context): - atype = opr.mode - if atype == "Identity": - context.set_blob_name(opr.out_vars[0], context.get_blob_name(opr.inp_vars[0])) - elif atype == "FUSE_ADD_RELU": - _arith(opr, "ADD", context) - bottom = [context.get_blob_name(opr.out_vars[0])] - top = [context.reset_blob_name(opr.out_vars[0], bottom[0] + ":relu")] - context.add_layer( - cp.LayerParameter(bottom=bottom, top=top, name=top[0], type="ReLU") - ) - elif atype == "FUSE_MUL_ADD3": - - mge_opr = opr._opr - mul_opr = ElemwiseOpr(mge_opr) - mul_opr.name = mul_opr.name + "_MUL" - mul_opr.add_inp_var(opr.inp_vars[0]) - mul_opr.add_inp_var(opr.inp_vars[1]) - mul_out = Tensor(opr.out_vars[0]._var, mge_opr) - mul_opr.add_out_var(mul_out) - _arith(mul_opr, "MUL", context) - - add_opr = ElemwiseOpr(mge_opr) - add_opr.name = add_opr.name + "_ADD" - add_opr.add_inp_var(mul_out) - add_opr.add_inp_var(opr.inp_vars[2]) - add_opr.add_out_var(opr.out_vars[0]) - _arith(add_opr, "ADD", context) - - elif atype == "RELU": - context.add_layer(_gen_layer(opr, "ReLU", context)) - elif atype == "TANH": - context.add_layer(_gen_layer(opr, "TanH", context)) - elif atype == "EXP": - context.add_layer(_gen_layer(opr, "Exp", context)) - elif atype == "SIGMOID": - context.add_layer(_gen_layer(opr, "Sigmoid", context)) - elif atype == "LOG": - context.add_layer(_gen_layer(opr, "Log", context)) - elif atype == "ABS": - context.add_layer(_gen_layer(opr, "AbsVal", context)) - elif atype in ["ADD", "SUB", "MUL", "TRUE_DIV"]: - _arith(opr, opr.mode, context) - elif atype == "POW": - power = opr.inp_vars[1].np_data - assert power.shape == (1,) - power_param = cp.PowerParameter(scale=1, shift=0, power=power[0]) - context.add_layer(_gen_layer(opr, "Power", context, power_param=power_param)) - elif atype == "MAX": - param = cp.EltwiseParameter(operation="MAX") - assert ( - opr.inp_vars[0].np_data is None and opr.inp_vars[1].np_data is None - ), "Caffe doesn't support elemwise MAX(tensor, const)" - context.add_layer( - _gen_layer(opr, "Eltwise", context, single_input=False, eltwise_param=param) - ) - else: - assert ( - False - ), "Elemwise op doesn't support mode {}, you can implement it in _eltwise".format( - opr.mode - ) +@_register_op(AbsOpr) +def _abs(tm_opr: OpBase, context): + context.add_layer(_gen_layer(tm_opr, "AbsVal", context)) -@_register_op( - ConvolutionForwardOpr, - ConvolutionBackwardDataOpr, - ConvForwardBiasOpr, - ConvolutionBackwardDataBiasOpr, -) +@_register_op(AddOpr, SubOpr, MulOpr, TrueDivOpr) +def _elemwise_arith(tm_opr: OpBase, context): + _arith(tm_opr, context) + + +@_register_op(IdentityOpr) +def _indentity(tm_opr: OpBase, context): + context.set_blob_name( + tm_opr.out_tensors[0], context.get_blob_name(tm_opr.inp_tensors[0]) + ) + + +@_register_op(Conv2dOpr, Deconv2dOpr) def _convolution(opr, context): - ph, pw = opr.ph, opr.pw - sh, sw = opr.sh, opr.sw - kh, kw = opr.kh, opr.kw - param_W = opr.param_W - group = opr.group - dilation_h, dilation_w = opr.dilation_h, opr.dilation_w + def expand(x): + if isinstance(x, (list, tuple)): + return x + elif isinstance(x, int): + return x, x + else: + raise TypeError( + "get error type! got {} expect int or tuple[int,..]".format(type(x)) + ) + + ph, pw = expand(opr.padding) + sh, sw = expand(opr.stride) + param_W = opr.inp_tensors[1].np_data + kh, kw = param_W.shape[-2:] + group = opr.groups + dilation_h, dilation_w = expand(opr.dilation) assert ( dilation_h == dilation_w ), "caffe accept one dilation, so dilation_h and dilation_w must equal" param_W = param_W.reshape((-1,) + param_W.shape[-3:]) - bias_term = opr.bias_term + bias_term = len(opr.inp_tensors) > 2 blobs = [ context.gen_blob_proto(param_W), ] param = cp.ConvolutionParameter( - num_output=opr.num_output, + num_output=opr.out_tensors[0].shape[1], pad_h=ph, pad_w=pw, stride_h=sh, @@ -645,16 +642,17 @@ def _convolution(opr, context): group=group, bias_term=bias_term, ) - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] - if isinstance(opr, ConvolutionBackwardDataOpr): + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] + if isinstance(opr, Deconv2dOpr): layer_type = "Deconvolution" - bottom = [context.get_blob_name(opr.inp_vars[1])] + context.set_blob_name(opr.inp_tensors[1], opr.inp_tensors[0].name) + bottom = [context.get_blob_name(opr.inp_tensors[1])] else: layer_type = "Convolution" - bottom = [context.get_blob_name(opr.inp_vars[0])] + bottom = [context.get_blob_name(opr.inp_tensors[0])] - if isinstance(opr, (ConvForwardBiasOpr, ConvolutionBackwardDataBiasOpr)): - blobs.append(context.gen_blob_proto(opr.param_B.reshape(-1,))) + if len(opr.inp_tensors) > 2: + blobs.append(context.gen_blob_proto(opr.inp_tensors[2].np_data.reshape(-1,))) assert bias_term == True else: assert bias_term == False @@ -663,7 +661,7 @@ def _convolution(opr, context): cp.LayerParameter( bottom=bottom, top=top, - name=opr.name, + name=opr.out_tensors[0].name, type=layer_type, blobs=blobs, convolution_param=param, @@ -671,34 +669,57 @@ def _convolution(opr, context): ) -@_register_op(PoolingForwardOpr) +@_register_op(AvgPool2dOpr, MaxPool2dOpr, AdaptiveAvgPool2dOpr) def _pooling2d(opr, context): - assert opr.mode in [ - "MAX", - "AVERAGE", - "AVERAGE_COUNT_EXCLUDE_PADDING", - ], "Pooling op doesn't support mode {}, you can implement it in _pooling2d".format( - opr.mode - ) - if opr.mode == "AVERAGE_COUNT_EXCLUDE_PADDING": - logger.warning( - "Caffe average pooling layer doesn't support 'COUNT_EXCLUDE_PADDING', you'd better set pooling mode to 'AVERAGE'" - ) - ph, pw = opr.ph, opr.pw - sh, sw = opr.sh, opr.sw - kh, kw = opr.kh, opr.kw - assert not None in opr.inp_vars[0].shape[2:4] - - ih, iw = opr.inp_vars[0].shape[2:4] - nh = ceil((ph * 2 + ih - kh + sh) / sh) - nw = ceil((pw * 2 + iw - kw + sw) / sw) - if ph > 0 and (nh - 1) * sh >= ih + ph: - nh = nh - 1 - if pw > 0 and (nw - 1) * sw >= iw + pw: - nw = nw - 1 + # assert opr.mode in [ + # "MAX", + # "AVERAGE", + # "AVERAGE_COUNT_EXCLUDE_PADDING", + # ], "Pooling op doesn't support mode {}, you can implement it in _pooling2d".format( + # opr.mode + # ) + def _unexpand(x): + if isinstance(x, Sequence): + return x[0], x[1] + elif isinstance(x, int): + return x, x + else: + raise TypeError( + "get error type! got {} expect int or tuple[int,..]".format(type(x)) + ) + + if not isinstance(opr, AdaptiveAvgPool2dOpr): + if opr.mode == "AVERAGE_COUNT_EXCLUDE_PADDING": + logger.warning( + "Caffe average pooling layer doesn't support 'COUNT_EXCLUDE_PADDING', you'd better set pooling mode to 'AVERAGE'" + ) + + ph, pw = _unexpand(opr.padding) + sh, sw = _unexpand(opr.stride) + kh, kw = _unexpand(opr.kernel_size) + assert not None in opr.inp_tensors[0].shape[2:4] + + ih, iw = opr.inp_tensors[0].shape[2:4] + nh = ceil((ph * 2 + ih - kh + sh) / sh) + nw = ceil((pw * 2 + iw - kw + sw) / sw) + if ph > 0 and (nh - 1) * sh >= ih + ph: + nh = nh - 1 + if pw > 0 and (nw - 1) * sw >= iw + pw: + nw = nw - 1 + elif isinstance(opr, AdaptiveAvgPool2dOpr): + oh, ow = _unexpand(opr.out_shape) + ih, iw = list(opr.inp_tensors[0].shape)[-2:] + ph, pw = 0, 0 + sh, sw = ih // oh, iw // ow + kh, kw = ih - (oh - 1) * sh, iw - (ow - 1) * sw + + if hasattr(opr, "mode"): + pool_mode = 0 if opr.mode == "MAX" else 1 + else: + pool_mode = 1 param = cp.PoolingParameter( - pool=0 if opr.mode == "MAX" else 1, + pool=pool_mode, pad_h=ph, pad_w=pw, stride_h=sh, @@ -707,54 +728,207 @@ def _pooling2d(opr, context): kernel_w=kw, ) - bottom = [context.get_blob_name(opr.inp_vars[0])] - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( - name=opr.name, type="Pooling", bottom=bottom, top=top, pooling_param=param + name=opr.out_tensors[0].name, + type="Pooling", + bottom=bottom, + top=top, + pooling_param=param, ) ) - if (nh - 1) * sh + kh > ph * 2 + ih: - logger.warning("Add extra 'slice layer' after Pooling Opr %s", opr.name) - param = cp.SliceParameter(axis=2, slice_point=[nh - 1]) - bottom = top[:1] - name = opr.name + context.gen_name - top = [name, context.gen_name] + if not isinstance(opr, AdaptiveAvgPool2dOpr): + if (nh - 1) * sh + kh > ph * 2 + ih: + logger.warning( + "Add extra 'slice layer' after Pooling Opr %s", opr.out_tensors[0].name + ) + param = cp.SliceParameter(axis=2, slice_point=[nh - 1]) + bottom = top[:1] + name = opr.out_tensors[0].name + context.gen_name + top = [name, context.gen_name] + context.add_layer( + cp.LayerParameter( + name=name, type="Slice", bottom=bottom, top=top, slice_param=param + ) + ) + context.add_layer(silence_blob(top[1])) + if (nw - 1) * sw + kw > pw * 2 + iw: + logger.warning( + "Add extra 'slice layer' after Pooling Opr %s", opr.out_tensors[0].name + ) + param = cp.SliceParameter(axis=3, slice_point=[nw - 1]) + bottom = top[:1] + name = opr.out_tensors[0].name + context.gen_name + top = [name, context.gen_name] + context.add_layer( + cp.LayerParameter( + name=name, type="Slice", bottom=bottom, top=top, slice_param=param + ) + ) + context.add_layer(silence_blob(top[1])) + context.reset_blob_name(opr.out_tensors[0], top[0]) + + +@_register_op(ConcatOpr) +def _concat(opr, context): + param = cp.ConcatParameter(axis=opr.axis) + context.add_layer( + _gen_layer(opr, "Concat", context, single_input=False, concat_param=param) + ) + + +@_register_op(SigmoidOpr, TanHOpr, LogOpr, ExpOpr) +def _elemwise_single_layer(opr, context): + context.add_layer(_gen_layer(opr, opr.name, context)) + + +@_register_op(MaxOpr) +def _elemwise_max(opr, context): + param = cp.EltwiseParameter(operation="MAX") + assert ( + opr.inp_tensors[0].np_data is None and opr.inp_tensors[1].np_data is None + ), "Caffe doesn't support elemwise MAX(tensor, const)" + context.add_layer( + _gen_layer(opr, "Eltwise", context, single_input=False, eltwise_param=param) + ) + + +@_register_op(MinOpr) +def _elemwise_min(opr, context): + param = cp.EltwiseParameter(operation="MIN") + assert ( + opr.inp_tensors[0].np_data is None and opr.inp_tensors[1].np_data is None + ), "Caffe doesn't support elemwise MIN(tensor, const)" + context.add_layer( + _gen_layer(opr, "Eltwise", context, single_input=False, eltwise_param=param) + ) + + +@_register_op(PowOpr) +def _pow(opr, context): + power = opr.inp_tensors[1].np_data + assert power.size == 1 # must with one power number + power_param = cp.PowerParameter(scale=1, shift=0, power=power.item()) + context.add_layer(_gen_layer(opr, "Power", context, power_param=power_param)) + + +@_register_op(ReduceOpr) +def _reduce(opr, context): + assert opr.mode in [ + "SUM", + "MAX", + "SUM_SQR", + "MEAN", + ], "Reduce op doesn't support mode {}, you can implement it in _reduce".format( + opr.mode + ) + if opr.mode in ("SUM", "SUM_SQR", "MEAN"): + mode = opr.mode + if opr.mode == "SUM_SQR": + mode = "SUMSQ" + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] + context.add_layer( + cp.LayerParameter( + name=opr.out_tensors[0].name, + type="Reduction", + bottom=bottom, + top=top, + reduction_param=cp.ReductionParameter(operation=mode, axis=opr.axis), + ) + ) + param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_tensors[0].shape)) + bottom = top + name = opr.out_tensors[0].name + context.gen_name + top = [context.reset_blob_name(opr.out_tensors[0], name)] context.add_layer( cp.LayerParameter( - name=name, type="Slice", bottom=bottom, top=top, slice_param=param + name=name, type="Reshape", bottom=bottom, top=top, reshape_param=param + ) + ) + + if opr.mode == "MAX": + logger.warning("Use 'slice/concat layers' to replace operator: Reduce Max") + bottom = [context.get_blob_name(opr.inp_tensors[0])] + name = opr.out_tensors[0].name + context.gen_name + top = [context.gen_name for _ in range(opr.inp_tensors[0].shape[opr.axis])] + context.add_layer( + cp.LayerParameter( + name=name, + type="Slice", + bottom=bottom, + top=top, + slice_param=cp.SliceParameter( + axis=opr.axis, + slice_point=list(range(1, opr.inp_tensors[0].shape[opr.axis])), + ), ) ) - context.add_layer(silence_blob(top[1])) - if (nw - 1) * sw + kw > pw * 2 + iw: - logger.warning("Add extra 'slice layer' after Pooling Opr %s", opr.name) - param = cp.SliceParameter(axis=3, slice_point=[nw - 1]) - bottom = top[:1] - name = opr.name + context.gen_name - top = [name, context.gen_name] + bottom = top + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( - name=name, type="Slice", bottom=bottom, top=top, slice_param=param + name=opr.out_tensors[0].name, + type="Eltwise", + bottom=bottom, + top=top, + eltwise_param=cp.EltwiseParameter(operation="MAX"), ) ) - context.add_layer(silence_blob(top[1])) - context.reset_blob_name(opr.out_vars[0], top[0]) + + if not opr.keep_dims: + param = cp.ReshapeParameter( + shape=cp.BlobShape(dim=opr.out_tensors[0].shape) + ) + bottom = top + name = opr.out_tensors[0].name + context.gen_name + top = [context.reset_blob_name(opr.out_tensors[0], name)] + context.add_layer( + cp.LayerParameter( + name=name, + type="Reshape", + bottom=bottom, + top=top, + reshape_param=param, + ) + ) + + +@_register_op(ReluOpr) +def _relu(tm_opr: OpBase, context): + context.add_layer(_gen_layer(tm_opr, "ReLU", context)) + + +@_register_op(SoftmaxOpr) +def _softmax(tm_opr: OpBase, context): + context.add_layer(_gen_layer(tm_opr, "Softmax", context)) @_register_op(ReshapeOpr) def _reshape(opr, context): - if isconst(opr.inp_vars[0]): - if opr.out_vars[0].np_data is None: - opr.out_vars[0].np_data = get_symvar_value(opr.out_vars[0]._var) + if isconst(opr.inp_tensors[0]): + # delete set out_tensor np data + if opr.out_tensors[0].np_data is None: + if hasattr(opr.out_tensors[0], "_var"): + opr.out_tensors[0].np_data = get_symvar_value(opr.out_tensors[0]._var) + else: + opr.out_tensors[0].np_data = opr.inp_tensors[0].np_data.reshape( + opr.out_shape + ) return - inp_shape = opr.input_shape - out_shape = opr.output_shape + inp_shape = opr.inp_tensors[0].shape + out_shape = tuple(opr.out_shape) if inp_shape == out_shape: - logger.info("Input shape and output shape of Opr %s are same, skip!", opr.name) - inp_blob = context.get_blob_name(opr.inp_vars[0]) - context.set_blob_name(opr.out_vars[0], inp_blob) + logger.info( + "Input shape and output shape of Opr %s are same, skip!", + opr.out_tensors[0].name, + ) + inp_blob = context.get_blob_name(opr.inp_tensors[0]) + context.set_blob_name(opr.out_tensors[0], inp_blob) elif ( all([inp_shape, out_shape]) and len(inp_shape) >= len(out_shape) @@ -762,8 +936,8 @@ def _reshape(opr, context): ): d = len(inp_shape) - len(out_shape) tmp_shape = out_shape + (1,) * d - bottom = [context.get_blob_name(opr.inp_vars[0])] - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] if inp_shape != tmp_shape: param = cp.ReshapeParameter(shape=cp.BlobShape(dim=tmp_shape)) tmp = [bottom[0] + "tmp"] @@ -771,26 +945,31 @@ def _reshape(opr, context): cp.LayerParameter( bottom=bottom, top=tmp, - name=opr.name + "tmp", + name=opr.out_tensors[0].name + "tmp", type="Reshape", reshape_param=param, ) ) bottom = tmp - param = cp.FlattenParameter(axis=len(out_shape) - 1, end_axis=-1) - context.add_layer( - cp.LayerParameter( - bottom=bottom, - top=top, - name=opr.name, - type="Flatten", - flatten_param=param, - ) + + logger.warning( + "trt don't support flatten after reshape, but caffe support! please do `export SUPPORT_CAFFE_TRT=True` or `SUPPORT_CAFFE_TRT=True your_command_line_xxxx` and do `unset SUPPORT_CAFFE_TRT` when running with caffe" ) + if "SUPPORT_CAFFE_TRT" not in os.environ: + param = cp.FlattenParameter(axis=len(out_shape) - 1, end_axis=-1) + context.add_layer( + cp.LayerParameter( + bottom=bottom, + top=top, + name=opr.out_tensors[0].name, + type="Flatten", + flatten_param=param, + ) + ) else: logger.warning( "NNIE doesn't support this reshape Opr %s, inp_shape %s, out_shape %s, NNIE reshape only support C/H/W, not N!", - opr.name, + opr.out_tensors[0].name, inp_shape, out_shape, ) @@ -800,39 +979,52 @@ def _reshape(opr, context): context.add_layer(_gen_layer(opr, "Reshape", context, reshape_param=param)) -@_register_op(ConcatOpr) -def _concat(opr, context): - param = cp.ConcatParameter(axis=opr.axis) +@_register_op(ConstantOpr) +def _constant(*_): + pass + + +@_register_op(FlattenOpr) +def _flatten_shape(opr, context): + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] + param = cp.FlattenParameter(axis=opr.start_axis, end_axis=opr.end_axis) context.add_layer( - _gen_layer(opr, "Concat", context, single_input=False, concat_param=param) + cp.LayerParameter( + bottom=bottom, + top=top, + name=opr.out_tensors[0].name, + type="Flatten", + flatten_param=param, + ) ) -@_register_op(MatrixMulOpr, FullyConnectedOpr) +@_register_op(MatMulOpr, LinearOpr) def _fully_connected(opr, context): - assert opr.inp_vars[1].np_data is not None - param_W = opr.inp_vars[1].np_data - assert not opr.transposeA + assert opr.inp_tensors[1].np_data is not None + param_W = opr.inp_tensors[1].np_data + assert not opr.transpose_a - if not opr.transposeB: + if not opr.transpose_b: param_W = param_W.T blobs = [context.gen_blob_proto(param_W)] bias_term = False - if isinstance(opr, FullyConnectedOpr): + if isinstance(opr, LinearOpr) and opr.has_bias: bias_term = True - blobs.append(context.gen_blob_proto(opr.param_B.reshape(-1,))) + blobs.append(context.gen_blob_proto(opr.inp_tensors[2].np_data.reshape(-1,))) param = cp.InnerProductParameter( - bias_term=bias_term, num_output=opr.out_vars[0].shape[1] + bias_term=bias_term, num_output=opr.out_tensors[0].shape[1] ) - bottom = [context.get_blob_name(opr.inp_vars[0])] - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( - name=opr.name, + name=opr.out_tensors[0].name, type="InnerProduct", bottom=bottom, top=top, @@ -842,121 +1034,29 @@ def _fully_connected(opr, context): ) -@_register_op(BatchNormForwardOpr) -def _batchnorm(opr, context): - inp = opr.inp_vars[0] - scale_ = opr.scale - bias_ = opr.bias - mean_ = opr.mean - var_ = opr.var - bottom = [ - context.get_blob_name(inp), - ] - tmp = [bottom[0] + context.gen_name] - top = [context.set_blob_name(opr.out_vars[opr.output_idx], opr.out_vars[0].name)] - bn_param = cp.BatchNormParameter(use_global_stats=True) - bn_blobs = [ - context.gen_blob_proto(mean_), - context.gen_blob_proto(var_), - context.gen_blob_proto(np.array([1])), - ] +@_register_op(SqueezeOpr) +def _squeeze(opr, context): + logger.warning("Use 'reshape layer' to replace operator: AxisAddRemove") + param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_tensors[0].shape)) + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( + name=opr.out_tensors[0].name, + type="Reshape", bottom=bottom, - top=tmp, - name=opr.name + "_bn", - type="BatchNorm", - batch_norm_param=bn_param, - blobs=bn_blobs, - ) - ) - - scale_param = cp.ScaleParameter(axis=1, num_axes=bias_.ndim) - scale_param.bias_term = True - scale_blobs = [context.gen_blob_proto(scale_), context.gen_blob_proto(bias_)] - context.add_layer( - cp.LayerParameter( - bottom=tmp, top=top, - name=opr.name + "_scale", - type="Scale", - scale_param=scale_param, - blobs=scale_blobs, + reshape_param=param, ) ) -@_register_op(ReduceOpr) -def _reduce(opr, context): - assert opr.mode in [ - "SUM", - "MAX", - "SUM_SQR", - "MEAN", - ], "Reduce op doesn't support mode {}, you can implement it in _reduce".format( - opr.mode - ) - if opr.mode in ("SUM", "SUM_SQR", "MEAN"): - mode = opr.mode - if opr.mode == "SUM_SQR": - mode = "SUMSQ" - bottom = [context.get_blob_name(opr.inp_vars[0])] - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] - context.add_layer( - cp.LayerParameter( - name=opr.name, - type="Reduction", - bottom=bottom, - top=top, - reduction_param=cp.ReductionParameter(operation=mode, axis=opr.axis), - ) - ) - param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_vars[0].shape)) - bottom = top - name = opr.name + context.gen_name - top = [context.reset_blob_name(opr.out_vars[0], name)] - context.add_layer( - cp.LayerParameter( - name=name, type="Reshape", bottom=bottom, top=top, reshape_param=param - ) - ) - - if opr.mode == "MAX": - logger.warning("Use 'slice/concat layers' to replace operator: Reduce Max") - bottom = [context.get_blob_name(opr.inp_vars[0])] - name = opr.name + context.gen_name - top = [context.gen_name for _ in range(opr.inp_vars[0].shape[opr.axis])] - context.add_layer( - cp.LayerParameter( - name=name, - type="Slice", - bottom=bottom, - top=top, - slice_param=cp.SliceParameter( - axis=opr.axis, - slice_point=list(range(1, opr.inp_vars[0].shape[opr.axis])), - ), - ) - ) - bottom = top - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] - context.add_layer( - cp.LayerParameter( - name=opr.name, - type="Eltwise", - bottom=bottom, - top=top, - eltwise_param=cp.EltwiseParameter(operation="MAX"), - ) - ) - - @_register_op(AxisAddRemoveOpr) -def axis_add_remove(opr, context): +def _axis_add_remove(opr, context): logger.warning("Use 'reshape layer' to replace operator: AxisAddRemove") - param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_vars[0].shape)) - bottom = [context.get_blob_name(opr.inp_vars[0])] - top = [context.set_blob_name(opr.out_vars[0], opr.out_vars[0].name)] + param = cp.ReshapeParameter(shape=cp.BlobShape(dim=opr.out_tensors[0].shape)) + bottom = [context.get_blob_name(opr.inp_tensors[0])] + top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)] context.add_layer( cp.LayerParameter( name=opr.name, type="Reshape", bottom=bottom, top=top, reshape_param=param @@ -965,22 +1065,82 @@ def axis_add_remove(opr, context): @_register_op(LeakyReluOpr) -def leaky_relu(opr, context): - param = cp.ReLUParameter(negative_slope=opr.negative_slope[0]) +def _leaky_relu(opr, context): + param = cp.ReLUParameter(negative_slope=opr.negative_slope) context.add_layer(_gen_layer(opr, "ReLU", context, relu_param=param)) +@_register_op(Relu6Opr) +def relu6(opr, context): + param = cp.ClipParameter(min=0.0, max=6.0) + context.add_layer(_gen_layer(opr, "Clip", context, clip_param=param)) + + +def add_3_relu6_div_6(opr, context, out): + inp = opr.inp_tensors[0] + const_3 = opr.inp_tensors[1] + const_6 = opr.inp_tensors[2] + fake_add_3_out = opr.inp_tensors[3] + fake_relu6_out = opr.inp_tensors[4] + split_add_op = AddOpr() + split_add_op.add_inp_tensors(inp) + split_add_op.add_inp_tensors(const_3) + split_add_op.add_out_tensors(fake_add_3_out) + _arith(split_add_op, context) + relu6_op = Relu6Opr() + relu6_op.add_inp_tensors(fake_add_3_out) + relu6_op.add_out_tensors(fake_relu6_out) + relu6(relu6_op, context) + true_div_op = TrueDivOpr() + true_div_op.add_inp_tensors(fake_relu6_out) + true_div_op.add_inp_tensors(const_6) + true_div_op.add_out_tensors(out) + _arith(true_div_op, context) + + +@_register_op(HardSigmoidOpr) +def hsigmoid(opr, context): + add_3_relu6_div_6(opr, context, opr.out_tensors[0]) + + +@_register_op(HardSwishOpr) +def hswish(opr, context): + inp = opr.inp_tensors[0] + fake_div6_out = opr.inp_tensors[5] + add_3_relu6_div_6(opr, context, fake_div6_out) + mul_op = MulOpr() + mul_op.add_inp_tensors(inp) + mul_op.add_inp_tensors(fake_div6_out) + mul_op.add_out_tensors(opr.out_tensors[0]) + _arith(mul_op, context) + + +@_register_op(SiLUOpr) +def silu(opr, context): + inp = opr.inp_tensors[0] + sigmoid_op = SigmoidOpr() + sigmoid_op.add_inp_tensors(inp) + fake_sigmoid_out = IRTensor(inp.name + "_sigmoid_out", inp.shape, inp.dtype,) + sigmoid_op.add_out_tensors(fake_sigmoid_out) + context.add_layer(_gen_layer(sigmoid_op, sigmoid_op.name, context)) + mul_op = MulOpr() + mul_op.add_inp_tensors(inp) + mul_op.add_inp_tensors(fake_sigmoid_out) + mul_op.add_out_tensors(opr.out_tensors[0]) + _arith(mul_op, context) + + @_register_op(TypeCvtOpr) -def typecvt(opr, context): - context.set_blob_name(opr.out_vars[0], context.get_blob_name(opr.inp_vars[0])) +def _typecvt(opr, context): + context.set_blob_name(opr.out_tensors[0], context.get_blob_name(opr.inp_tensors[0])) @_register_op(BroadcastOpr) -def broadcast(opr, context): - input = opr.inp_vars[0] +def _broadcast(opr, context): + input = opr.inp_tensors[0] inp_ndim = input.ndim a_shape = input.shape - b_shape = opr.inp_vars[1].np_data + b_shape = opr.inp_tensors[1].np_data b_ndim = len(b_shape) assert inp_ndim <= b_ndim bottom = [context.get_blob_name(input)] @@ -993,7 +1153,7 @@ def broadcast(opr, context): cp.LayerParameter( bottom=bottom, top=top, - name=opr.name + context.gen_name, + name=opr.out_tensors[0].name + context.gen_name, type="Reshape", reshape_param=param, ) @@ -1002,7 +1162,7 @@ def broadcast(opr, context): for i in range(b_ndim): shpA, shpB = a_shape[i], b_shape[i] assert shpA in (shpB, 1) - name = opr.name + context.gen_name + name = opr.out_tensors[0].name + context.gen_name top = [name] context.add_layer( cp.LayerParameter( @@ -1014,4 +1174,52 @@ def broadcast(opr, context): ) ) bottom = top - context.set_blob_name(opr.out_vars[0], bottom[0]) + context.set_blob_name(opr.out_tensors[0], bottom[0]) + + +@_register_op(BatchNormalizationOpr) +def _batchnorm(opr, context): + inp = opr.inp_tensors[0] + scale_ = opr.inp_tensors[1].np_data.squeeze() + bias_ = opr.inp_tensors[2].np_data.squeeze() + mean_ = opr.inp_tensors[3].np_data.squeeze() + var_ = opr.inp_tensors[4].np_data.squeeze() + bottom = [ + context.get_blob_name(inp), + ] + tmp = [bottom[0] + context.gen_name] + top = [ + context.set_blob_name( + opr.out_tensors[opr.output_idx], opr.out_tensors[opr.output_idx].name + ) + ] + bn_param = cp.BatchNormParameter(use_global_stats=True) + bn_blobs = [ + context.gen_blob_proto(mean_), + context.gen_blob_proto(var_), + context.gen_blob_proto(np.array([1])), + ] + context.add_layer( + cp.LayerParameter( + bottom=bottom, + top=tmp, + name=opr.out_tensors[opr.output_idx].name + "_bn", + type="BatchNorm", + batch_norm_param=bn_param, + blobs=bn_blobs, + ) + ) + + scale_param = cp.ScaleParameter(axis=1, num_axes=bias_.ndim) + scale_param.bias_term = True + scale_blobs = [context.gen_blob_proto(scale_), context.gen_blob_proto(bias_)] + context.add_layer( + cp.LayerParameter( + bottom=tmp, + top=top, + name=opr.out_tensors[opr.output_idx].name + "_scale", + type="Scale", + scale_param=scale_param, + blobs=scale_blobs, + ) + ) diff --git a/mgeconvert/caffe_converter/init.sh b/mgeconvert/backend/ir_to_caffe/init.sh similarity index 88% rename from mgeconvert/caffe_converter/init.sh rename to mgeconvert/backend/ir_to_caffe/init.sh index 4c1ef26..76721d3 100755 --- a/mgeconvert/caffe_converter/init.sh +++ b/mgeconvert/backend/ir_to_caffe/init.sh @@ -1,6 +1,6 @@ #!/bin/bash -e -python3 -m pip install protobuf --user +python3 -m pip install --no-binary=protobuf protobuf==3.8.0 --user hash wget || (echo "please install wget package" && exit -1) hash protoc || (echo "please install protobuf-compiler package" && exit -1) diff --git a/mgeconvert/backend/ir_to_onnx/__init__.py b/mgeconvert/backend/ir_to_onnx/__init__.py new file mode 100644 index 0000000..114ee8e --- /dev/null +++ b/mgeconvert/backend/ir_to_onnx/__init__.py @@ -0,0 +1 @@ +from .onnx_converter import OnnxConverter diff --git a/mgeconvert/onnx_converter/init.sh b/mgeconvert/backend/ir_to_onnx/init.sh similarity index 100% rename from mgeconvert/onnx_converter/init.sh rename to mgeconvert/backend/ir_to_onnx/init.sh diff --git a/mgeconvert/onnx_converter/onnx_converter.py b/mgeconvert/backend/ir_to_onnx/onnx_converter.py similarity index 62% rename from mgeconvert/onnx_converter/onnx_converter.py rename to mgeconvert/backend/ir_to_onnx/onnx_converter.py index e8a8c4f..b2533c1 100644 --- a/mgeconvert/onnx_converter/onnx_converter.py +++ b/mgeconvert/backend/ir_to_onnx/onnx_converter.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,37 +5,29 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from typing import List - import megengine as mge import onnx.checker import onnx.helper import onnx.numpy_helper from onnx import optimizer -from ..mge_context import TopologyNetwork, TransformerRule, optimize_for_conversion -from ..mge_context.mge_utils import get_symvar_value -from .onnx_op import MGE2ONNX, mge2onnx_dtype_mapping, set_opset_version +from ...frontend.mge_to_ir.mge_utils import get_symvar_value +from .onnx_op import ( + MGE2ONNX, + _add_input_tensors, + mge2onnx_dtype_mapping, + set_opset_version, +) class OnnxConverter: - transformer_options: List[TransformerRule] = [] - - def __init__( - self, toponet, transform_options=None, opset_version=8, graph_name="graph" - ): - assert isinstance( - toponet, TopologyNetwork - ), "net must be instance of TopologyNetwork" - self.net = toponet + def __init__(self, net, opset_version=8, graph_name="graph"): + self.net = net assert 7 <= opset_version <= 12, "opset {} are not supported yet".format( opset_version ) self.graph_name = graph_name self.opset_version = opset_version - if transform_options is not None: - self.transformer_options = transform_options - optimize_for_conversion(self.net, self.transformer_options) def convert(self): inputs = [] @@ -47,8 +38,8 @@ def convert(self): set_opset_version(self.opset_version) def need_convert(opr): - is_const = [data.np_data is not None for data in opr.inp_vars] - return not all(is_const) or len(opr.inp_vars) == 0 + is_const = [data.np_data is not None for data in opr.inp_tensors] + return not all(is_const) or len(opr.inp_tensors) == 0 def deduplication(inputs): names = [] @@ -59,10 +50,13 @@ def deduplication(inputs): names.append(i.name) return results + _, tensor_sources, _ = _add_input_tensors(self.net.graph_inputs) + inputs.extend(tensor_sources) + for opr in self.net.all_oprs: if not need_convert(opr): - for tensor in opr.out_vars: - if tensor.np_data is None: + for tensor in opr.out_tensors: + if hasattr(tensor, "_var"): tensor.np_data = get_symvar_value(tensor._var) continue converter_cls = MGE2ONNX.get(type(opr), None) @@ -83,7 +77,7 @@ def deduplication(inputs): unsupported_oprs ) - for output in self.net.output_vars: + for output in self.net.graph_outputs: def _get_onnx_dtype(output): return mge2onnx_dtype_mapping[output.dtype] @@ -111,29 +105,3 @@ def _get_onnx_dtype(output): ] model = optimizer.optimize(model, passes) return model - - -def convert_to_onnx( - mge_fpath, output="out.onnx", *, graph_name="graph", opset=8, outspec=None -): - """ - Convert megengine model to ONNX, - and save the ONNX model to file `output`. - - :param mge_fpath: the file path of megengine model. - :type fpath: str - :param output: the filename used for the saved model. - :type output: str - :param graph_name: the name of the ONNX graph. - :type graph_name: str - :param opset: opset version of ONNX model. - :type opset: int - """ - assert isinstance(mge_fpath, str), "mge_fpath must be string" - net = TopologyNetwork(mge_fpath, prune_reshape=True, outspec=outspec) - converter = OnnxConverter(net, None, opset, graph_name) - model = converter.convert() - - assert isinstance(output, str), "onnx_fpath must be string" - with open(output, "wb") as fout: - fout.write(model.SerializeToString()) diff --git a/mgeconvert/backend/ir_to_onnx/onnx_op.py b/mgeconvert/backend/ir_to_onnx/onnx_op.py new file mode 100644 index 0000000..df38c83 --- /dev/null +++ b/mgeconvert/backend/ir_to_onnx/onnx_op.py @@ -0,0 +1,1245 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +import numpy as np +import onnx + +from ...converter_ir.ir_op import ( + AbsOpr, + AdaptiveAvgPool2dOpr, + AddOpr, + AvgPool2dOpr, + AxisAddRemoveOpr, + BatchNormalizationOpr, + BroadcastOpr, + CeilOpr, + ConcatOpr, + ConstantOpr, + Conv2dOpr, + ConvolutionBackwardFilterOpr, + Deconv2dOpr, + DropoutOpr, + ExpOpr, + FlattenOpr, + FloorOpr, + FuseMulAdd3Opr, + GetSubTensorOpr, + GetVarShapeOpr, + HardSigmoidOpr, + HardSwishOpr, + IdentityOpr, + LinearOpr, + LinspaceOpr, + LogOpr, + MatMulOpr, + MaxOpr, + MaxPool2dOpr, + MinOpr, + MulOpr, + MultipleDeviceTensorHolderOpr, + PowOpr, + ReduceOpr, + Relu6Opr, + ReluOpr, + RepeatOpr, + ReshapeOpr, + SharedDeviceTensorOpr, + SigmoidOpr, + SiLUOpr, + SoftmaxOpr, + SqueezeOpr, + SubOpr, + TanHOpr, + TransposeOpr, + TrueDivOpr, + TypeCvtOpr, +) +from ...frontend.mge_to_ir.mge_utils import get_symvar_value + +mge2onnx_dtype_mapping = { + # pylint: disable=no-member + np.float32: onnx.TensorProto.FLOAT, + np.float16: onnx.TensorProto.FLOAT16, + np.int8: onnx.TensorProto.INT8, + np.int16: onnx.TensorProto.INT16, + np.int32: onnx.TensorProto.INT32, + np.int64: onnx.TensorProto.INT64, + np.uint8: onnx.TensorProto.UINT8, +} + +MGE2ONNX = {} + + +def _register_op(*oprs): + def callback(impl): + for opr in oprs: + MGE2ONNX[opr] = impl + return impl + + return callback + + +opset_version = 8 + + +def set_opset_version(version): + global opset_version # pylint: disable=W0603 + opset_version = version + + +def expand(x): + if isinstance(x, (list, tuple)): + return x + elif isinstance(x, int): + return x, x + else: + raise TypeError( + "get error type! got {} expect int or tuple[int,..]".format(type(x)) + ) + + +mge2onnx_dtype_mapping = { + # pylint: disable=no-member + np.float32: onnx.TensorProto.FLOAT, + np.float16: onnx.TensorProto.FLOAT16, + np.int8: onnx.TensorProto.INT8, + np.int16: onnx.TensorProto.INT16, + np.int32: onnx.TensorProto.INT32, + np.int64: onnx.TensorProto.INT64, + np.uint8: onnx.TensorProto.UINT8, +} + + +def _add_input_tensors(inputs): + inp_tensor_list = [ + onnx.helper.make_tensor_value_info( + tensor.name, mge2onnx_dtype_mapping[tensor.dtype], tensor.shape + ) + for tensor in inputs + ] + return [], inp_tensor_list, [] + + +class OperatorBaseConverter: + + __opr_type__ = "OperatorBaseConverter" + + def __init__(self, opr): + """ + :param opr: the operator that converter converts. + :type opr: subclass of :class:`.MgeOpr` + """ + self._opr = opr + self._net_sources = [] + self._parameters = [] + + def _get_inputs(self, exclude_idx=None): + """ + Returns the names of inputs of onnx operator. + """ + if exclude_idx is None: + exclude_idx = [] + for idx, inp in enumerate(self._opr.inp_tensors): + if idx not in exclude_idx: + if self._opr.inp_tensors[idx].np_data is not None: + inp_tensor = onnx.helper.make_tensor_value_info( + inp.name, mge2onnx_dtype_mapping[inp.dtype], inp.shape + ) + param = onnx.numpy_helper.from_array(inp.np_data, inp.name) + self._net_sources.append(inp_tensor) + self._parameters.append(param) + + return [tensor.name for tensor in self._opr.inp_tensors] + + def _get_outputs(self): + """ + Returns the names of outputs of onnx operator. + """ + return [tensor.name for tensor in self._opr.out_tensors] + + def _get_attrs(self): + """ + Returns extra attributes needed by :method:`.convert` + """ + return {} + + def convert(self): + """ + Convert owning operator to onnx operator. Could be override by + subclass. + + Returns tuple (nodes, net_sources, parameters) + """ + nodes = [ + onnx.helper.make_node( + self.__opr_type__, + self._get_inputs(), + self._get_outputs(), + name=self._opr.out_tensors[0].name, + **self._get_attrs(), + ) + ] + return nodes, self._net_sources, self._parameters + + +@_register_op( + AddOpr, + SubOpr, + TrueDivOpr, + MulOpr, + MinOpr, + ReluOpr, + ExpOpr, + TanHOpr, + SigmoidOpr, + AbsOpr, + LogOpr, + FloorOpr, + CeilOpr, + PowOpr, + MaxOpr, + FuseMulAdd3Opr, + IdentityOpr, + SiLUOpr, +) +class ElemwiseConverter(OperatorBaseConverter): + + support_op_map = { + AddOpr: "Add", + SubOpr: "Sub", + TrueDivOpr: "Div", + MulOpr: "Mul", + ReluOpr: "Relu", + Relu6Opr: "FUSE_RELU6", + ExpOpr: "Exp", + TanHOpr: "Tanh", + SiLUOpr: "SiLU", + SigmoidOpr: "Sigmoid", + AbsOpr: "Abs", + LogOpr: "Log", + FloorOpr: "Floor", + CeilOpr: "Ceil", + PowOpr: "Pow", + MaxOpr: "Max", + MinOpr: "Min", + IdentityOpr: "Identity", + } + + def __init__(self, opr): + super().__init__(opr) + assert isinstance( + opr, tuple(self.support_op_map.keys()) + ), "Elemwise op doesn't support mode {}, you can implement it in ElemwiseConverter".format( + type(opr) + ) + op_type = self.support_op_map[type(opr)] + self.__opr_type__ = op_type + + def convert(self): + if self.__opr_type__ == "SiLU": + opr = self._opr + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = [] + neg_x = inputs[0] + "_negtive" + neg_node = onnx.helper.make_node("Neg", [inputs[0]], [neg_x]) + nodes.append(neg_node) + exp = neg_x + "_exp" + exp_node = onnx.helper.make_node("Exp", [neg_x], [exp]) + nodes.append(exp_node) + const_1 = inputs[0] + "_const_1" + const_1_node = onnx.helper.make_node( + "Constant", + [], + [const_1], + value=onnx.helper.make_tensor( + const_1, mge2onnx_dtype_mapping[opr.inp_tensors[0].dtype], [], [1.0] + ), + ) + nodes.append(const_1_node) + add = exp + "_add_const_1" + add_node = onnx.helper.make_node("Add", [exp, const_1], [add]) + nodes.append(add_node) + div_node = onnx.helper.make_node("Div", [inputs[0], add], outputs) + nodes.append(div_node) + return nodes, self._net_sources, self._parameters + else: + return super().convert() + + +@_register_op(MultipleDeviceTensorHolderOpr, SharedDeviceTensorOpr, LinspaceOpr) +class IgnoredOperatorConverter(OperatorBaseConverter): + def convert(self): + return [], [], [] + + +@_register_op(GetVarShapeOpr) +class GetVarShapeConverter(OperatorBaseConverter): + __opr_type__ = "Shape" + + def convert(self): + shape = self._opr.out_tensors[0] + shape.np_data = get_symvar_value(shape._var).astype(np.int64) + shape.dtype = np.int64 + return [], [], [] + + +@_register_op(SoftmaxOpr) +class SoftmaxConverter(OperatorBaseConverter): + __opr_type__ = "Softmax" + + def _get_attrs(self): + if self._opr.axis is not None: + return {"axis": self._opr.axis} + else: + return {} + + def convert(self): + opr = self._opr + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = [] + assert opr.inp_tensors[0].ndim == 2, "ONNX Softmax only support dim=2" + offset_name = inputs[0] + "_max_offset" + offset = onnx.helper.make_node( + "ReduceMax", + inputs=[inputs[0]], + outputs=[offset_name], + axes=[opr.axis], + keepdims=True, + ) + nodes.append(offset) + sub_name = inputs[0] + "_sub_offset" + sub = onnx.helper.make_node( + "Sub", inputs=[inputs[0], offset_name], outputs=[sub_name], + ) + nodes.append(sub) + softmax = onnx.helper.make_node( + "Softmax", inputs=[sub_name], outputs=[outputs[0]], **self._get_attrs(), + ) + nodes.append(softmax) + return nodes, self._net_sources, self._parameters + + +@_register_op(GetSubTensorOpr) +class SubtensorConverter(OperatorBaseConverter): + + __opr_type__ = "Slice" + + def slice_version_1(self, starts, ends, axes, _, inputs, outputs): + attr = {"axes": axes, "ends": ends, "starts": starts} + slice_op = onnx.helper.make_node("Slice", inputs, outputs, **attr) + return slice_op, [], [] + + def slice_version_10(self, starts, ends, axes, steps, inputs, outputs): + op_name = self._opr.out_tensors[0].name + inputs = inputs + [ + op_name + "_begin", + op_name + "_end", + op_name + "_axis", + op_name + "_step", + ] + begin = onnx.helper.make_tensor_value_info( + inputs[1], mge2onnx_dtype_mapping[np.int32], starts.shape + ) + end = onnx.helper.make_tensor_value_info( + inputs[2], mge2onnx_dtype_mapping[np.int32], ends.shape + ) + axis = onnx.helper.make_tensor_value_info( + inputs[3], mge2onnx_dtype_mapping[np.int32], axes.shape + ) + step = onnx.helper.make_tensor_value_info( + inputs[4], mge2onnx_dtype_mapping[np.int32], steps.shape + ) + net_sources = [begin, end, axis, step] + parameters = [ + onnx.numpy_helper.from_array(starts, inputs[1]), + onnx.numpy_helper.from_array(ends, inputs[2]), + onnx.numpy_helper.from_array(axes, inputs[3]), + onnx.numpy_helper.from_array(steps, inputs[4]), + ] + Slice = onnx.helper.make_node("Slice", inputs, outputs) + return Slice, net_sources, parameters + + def convert(self): + opr = self._opr + squeeze_axis = opr.squeeze_axis + begin_param = np.array(opr.begin_params, dtype=np.int32) + end_param = np.array(opr.end_params, dtype=np.int32) + step_param = np.array(opr.step_params, dtype=np.int32) + axis_param = np.array(opr.axis, dtype=np.int32) + + inputs = [self._get_inputs(exclude_idx=list(range(1, len(opr.inp_tensors))))[0]] + outputs = self._get_outputs() + slice_outputs = [ + outputs[0] if len(squeeze_axis) == 0 else outputs[0] + "tmp@onnx" + ] + slice_op = None + slice_net_sources = [] + slice_param = [] + if opset_version < 10: + slice_op, slice_net_sources, slice_param = self.slice_version_1( + begin_param, end_param, axis_param, step_param, inputs, slice_outputs + ) + else: + slice_op, slice_net_sources, slice_param = self.slice_version_10( + begin_param, end_param, axis_param, step_param, inputs, slice_outputs + ) + nodes = [] + self._parameters.extend(slice_param) + self._net_sources.extend(slice_net_sources) + nodes.append(slice_op) + if len(squeeze_axis) > 0: + Squeeze = onnx.helper.make_node( + "Squeeze", slice_outputs, outputs, axes=squeeze_axis + ) + nodes.append(Squeeze) + + return (nodes, self._net_sources, self._parameters) + + +@_register_op(TransposeOpr) +class DimshuffleConverter(OperatorBaseConverter): + + __opr_type__ = "Transpose" + + def _get_attrs(self): + return {"perm": list(self._opr.pattern)} + + +@_register_op(MatMulOpr, LinearOpr) +class MatrixMulConvert(OperatorBaseConverter): + def convert(self): + opr = self._opr + nodes = [] + const_0 = opr.out_tensors[0].name + "_const_0_onnx" + const_0_tensor = onnx.helper.make_tensor_value_info( + const_0, mge2onnx_dtype_mapping[np.float32], [1] + ) + const_0_param = onnx.numpy_helper.from_array( + np.array([0]).astype("float32"), const_0 + ) + self._net_sources.append(const_0_tensor) + self._parameters.append(const_0_param) + inputs = self._get_inputs() + outputs = self._get_outputs() + if isinstance(opr, LinearOpr) and opr.has_bias: + temp_out = inputs[0] + "_mul" + inputs[1] + gemm = onnx.helper.make_node( + "Gemm", + [inputs[0], inputs[1], const_0], + [temp_out], + alpha=1.0, + beta=0.0, + transA=opr.transpose_a, + transB=opr.transpose_b, + ) + nodes.append(gemm) + add_bias = onnx.helper.make_node( + "Add", inputs=[temp_out, inputs[2]], outputs=[outputs[0]], + ) + nodes.append(add_bias) + else: + gemm = onnx.helper.make_node( + "Gemm", + [inputs[0], inputs[1], const_0], + [outputs[0]], + alpha=1.0, + beta=0.0, + transA=opr.transpose_a, + transB=opr.transpose_b, + ) + nodes.append(gemm) + + return ( + nodes, + self._net_sources, + self._parameters, + ) + + +@_register_op(ReshapeOpr) +class ReshapeConverter(OperatorBaseConverter): + + __opr_type__ = "Reshape" + + def convert(self): + inputs = self._get_inputs() + outputs = self._get_outputs() + shape_tensor_name = self._opr.out_tensors[0].name + "_shape_onnx" + shape_tensor = onnx.helper.make_tensor_value_info( + shape_tensor_name, + mge2onnx_dtype_mapping[np.int64], + (len(self._opr.out_shape),), + ) + shape_param = onnx.numpy_helper.from_array( + np.array(self._opr.out_shape, dtype=np.int64), shape_tensor_name + ) + reshape = onnx.helper.make_node( + "Reshape", [inputs[0], shape_tensor_name], outputs + ) + return ( + [reshape], + self._net_sources + [shape_tensor], + self._parameters + [shape_param], + ) + + +@_register_op(AxisAddRemoveOpr) +class AxisAddRemoveConverter(OperatorBaseConverter): + def convert(self): + inputs = self._get_inputs() + outputs = self._get_outputs() + add_axis = [] + remove_axis = [] + for desc in self._opr.desc: + if desc["method"] == 0: + add_axis.append(desc["axisnum"]) + else: + remove_axis.append(desc["axisnum"]) + + if len(add_axis) > 0 and len(remove_axis) > 0: + assert ( + False + ), "AsixAddRemove converter doesn't support add and remove axises concurrently" + + elif len(add_axis) > 0: + unsqueeze = onnx.helper.make_node( + "Unsqueeze", inputs, outputs, axes=add_axis + ) + ret = [unsqueeze] + elif len(remove_axis) > 0: + squeeze = onnx.helper.make_node( + "Squeeze", inputs, outputs, axes=remove_axis + ) + ret = [squeeze] + else: + ret = [] + return ret, self._net_sources, self._parameters + + +@_register_op(Conv2dOpr, Deconv2dOpr) +class Conv2DConverter(OperatorBaseConverter): + + __opr_type__ = "Conv" + + def _get_attrs(self): + opr = self._opr + ph, pw = expand(opr.padding) + sh, sw = expand(opr.stride) + param_W = opr.inp_tensors[1].shape + kh, kw = param_W[-2:] + group = opr.groups if opr.groups is not None else 1 + dilation_h, dilation_w = expand(opr.dilation) + return { + "kernel_shape": [kh, kw], + "pads": [ph, pw, ph, pw], + "strides": [sh, sw], + "dilations": [dilation_h, dilation_w], + "group": group if group is not None else 1, + } + + def convert(self): + opr = self._opr + attrs = self._get_attrs() + nodes = [] + exclude_idx = [0] if attrs["group"] != 1 else [] + inputs = self._get_inputs(exclude_idx) + outputs = self._get_outputs() + if attrs["group"] != 1: + flt_shape = self._opr.inp_tensors[1].shape + flt_shape = [ + flt_shape[0] * flt_shape[1], + flt_shape[2], + flt_shape[3], + flt_shape[4], + ] + + if opr.inp_tensors[1].np_data is not None: + inputs[1] = opr.out_tensors[0].name + "_filter_reshape_onnx" + flt = opr.inp_tensors[1].np_data + flt_data = flt.reshape(flt_shape) + flt_tensor = onnx.helper.make_tensor_value_info( + inputs[1], mge2onnx_dtype_mapping[flt.dtype.type], flt_shape + ) + flt_param = onnx.numpy_helper.from_array(flt_data, inputs[1]) + self._net_sources.append(flt_tensor) + self._parameters.append(flt_param) + else: + reshape_inputs = [inputs[1], opr.out_tensors[0].name + "shape_onnx"] + shape_tensor = onnx.helper.make_tensor_value_info( + reshape_inputs[1], + mge2onnx_dtype_mapping[np.int64], + (len(flt_shape),), + ) + shape_param = onnx.numpy_helper.from_array( + np.array(flt_shape, dtype="int64"), reshape_inputs[1] + ) + self._net_sources.append(shape_tensor) + self._parameters.append(shape_param) + reshape = onnx.helper.make_node( + "Reshape", + reshape_inputs, + [opr.out_tensors[0].name + "_filter_reshape_onnx"], + ) + inputs[1] = opr.out_tensors[0].name + "_filter_reshape_onnx" + nodes.append(reshape) + onnx_op = "Conv" + if isinstance(self._opr, Deconv2dOpr): + onnx_op = "ConvTranspose" + conv2d = onnx.helper.make_node(onnx_op, inputs, [outputs[0]], **attrs) + nodes.append(conv2d) + return (nodes, self._net_sources, self._parameters) + + +@_register_op(ConvolutionBackwardFilterOpr) +class Conv2DBackwardFilterConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + # src, grad_out, weight = self._opr.inp_vars + + nodes = [] + inputs = self._get_inputs() + outputs = self._get_outputs() + + # Tile + # grad_out: (no, co, ho, wo) -> (no, co x ci / group, ho, wo) + grad_out_tile_in = outputs[0] + "_grad_out_tile_in" + grad_out_tile_source = onnx.helper.make_tensor_value_info( + grad_out_tile_in, mge2onnx_dtype_mapping[np.int64], (4,) + ) + grad_out_tile_param = onnx.numpy_helper.from_array( + np.array([1, opr.src_shape[1] // opr.group, 1, 1]), grad_out_tile_in + ) + self._net_sources.append(grad_out_tile_source) + self._parameters.append(grad_out_tile_param) + grad_out_tile_out = outputs[0] + "_grad_out_tile_out" + grad_out_tile = onnx.helper.make_node( + "Tile", [inputs[1], grad_out_tile_in], [grad_out_tile_out] + ) + nodes.append(grad_out_tile) + + # Reshape + # grad_out: (no, co x ci / group, ho, wo) -> (no x co x ci / group, 1, ho, wo) + grad_out_reshape_in = outputs[0] + "_grad_out_reshape_in" + grad_out_reshape_source = onnx.helper.make_tensor_value_info( + grad_out_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) + ) + grad_out_reshape_param = onnx.numpy_helper.from_array( + np.array( + [ + opr.grad_out_shape[0] + * opr.grad_out_shape[1] + * opr.src_shape[1] + // opr.group, + 1, + opr.grad_out_shape[2], + opr.grad_out_shape[3], + ] + ), + grad_out_reshape_in, + ) + self._net_sources.append(grad_out_reshape_source) + self._parameters.append(grad_out_reshape_param) + grad_out_reshape_out = outputs[0] + "_grad_out_reshape_out" + reshape = onnx.helper.make_node( + "Reshape", [grad_out_tile_out, grad_out_reshape_in], [grad_out_reshape_out] + ) + nodes.append(reshape) + + # Reshape + # src: (ni, ci, hi, wi) -> (1, ni x ci, hi, wi) + src_reshape_in = outputs[0] + "_src_reshape_in" + src_reshape_source = onnx.helper.make_tensor_value_info( + src_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) + ) + src_reshape_param = onnx.numpy_helper.from_array( + np.array( + [ + 1, + opr.src_shape[0] * opr.src_shape[1], + opr.src_shape[2], + opr.src_shape[3], + ] + ), + src_reshape_in, + ) + self._net_sources.append(src_reshape_source) + self._parameters.append(src_reshape_param) + src_reshape_out = outputs[0] + "_src_reshape_out" + reshape = onnx.helper.make_node( + "Reshape", [inputs[0], src_reshape_in], [src_reshape_out] + ) + nodes.append(reshape) + + # Conv: + # group = ni * ci + # src(1, ni x ci, hi, wi) + grad_out(no x co x ci / group, 1, ho, wo) + # -> grad_weight(1, no x co x ci / group, ?, ?) + grad_weight = outputs[0] + "_grad_weight" + grad_weight_conv = onnx.helper.make_node( + "Conv", + [src_reshape_out, grad_out_reshape_out], + [grad_weight], + kernel_shape=[opr.grad_out_shape[2], opr.grad_out_shape[3]], + strides=[opr.dilation[0], opr.dilation[1]], + pads=[opr.padding[0], opr.padding[1], opr.padding[0], opr.padding[1]], + dilations=[opr.stride[0], opr.stride[1]], + group=opr.src_shape[1] * opr.src_shape[0], + ) + nodes.append(grad_weight_conv) + + # Slice + # grad_weight: (1, no x co x ci // group, ?, ?) -> (1, no x co x ci // group, kh, kw) + grad_weight_slice_out = outputs[0] + "_grad_weight_slice_out" + grad_weight_slice = onnx.helper.make_node( + "Slice", + [grad_weight], + [grad_weight_slice_out], + axes=[2, 3], + starts=[0, 0], + ends=[opr.kernel_shape[0], opr.kernel_shape[1]], + ) + nodes.append(grad_weight_slice) + + # Reshape + # grad_weight: (1, no x co x ci // group, kh, kw) -> (no, co x ci // group, kh, kw) + grad_weight_reshape_in = outputs[0] + "_grad_weight_reshape_in" + grad_weight_reshape_source = onnx.helper.make_tensor_value_info( + grad_weight_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) + ) + grad_weight_reshape_param = onnx.numpy_helper.from_array( + np.array( + [ + opr.grad_out_shape[0], + opr.grad_out_shape[1] * opr.src_shape[1] // opr.group, + opr.kernel_shape[0], + opr.kernel_shape[1], + ] + ), + grad_weight_reshape_in, + ) + self._net_sources.append(grad_weight_reshape_source) + self._parameters.append(grad_weight_reshape_param) + grad_weight_reshape_out = outputs[0] + "_grad_weight_reshape_out" + reshape = onnx.helper.make_node( + "Reshape", + [grad_weight_slice_out, grad_weight_reshape_in], + [grad_weight_reshape_out], + ) + nodes.append(reshape) + + # ReduceSum + # grad_weight: (no, co x ci // group, kh, kw) -> (1, co x ci // goup, kh, kw) + grad_weight_reduce_out = outputs[0] + "_grad_weight_reduce_out" + grad_weight_reduce = onnx.helper.make_node( + "ReduceSum", [grad_weight_reshape_out], [grad_weight_reduce_out], axes=[0], + ) + nodes.append(grad_weight_reduce) + + # Reshape + # grad_weight: (1, co x ci // group, kh, kw) -> (ci // group, co, kh, kw) + grad_weight_reshape2_in = outputs[0] + "_grad_weight_reshape2_in" + grad_weight_reshape2_source = onnx.helper.make_tensor_value_info( + grad_weight_reshape2_in, mge2onnx_dtype_mapping[np.int64], (4,) + ) + grad_weight_reshape2_param = onnx.numpy_helper.from_array( + np.array( + [ + opr.src_shape[1] // opr.group, + opr.grad_out_shape[1], + opr.kernel_shape[0], + opr.kernel_shape[1], + ] + ), + grad_weight_reshape2_in, + ) + self._net_sources.append(grad_weight_reshape2_source) + self._parameters.append(grad_weight_reshape2_param) + grad_weight_reshape2_out = outputs[0] + "_grad_weight_reshape2_out" + reshape = onnx.helper.make_node( + "Reshape", + [grad_weight_reduce_out, grad_weight_reshape2_in], + [grad_weight_reshape2_out], + ) + nodes.append(reshape) + + # Transpose + grad_weight_transpose_out = outputs[0] + transpose = onnx.helper.make_node( + "Transpose", + [grad_weight_reshape2_out], + [grad_weight_transpose_out], + perm=[1, 0, 2, 3], + ) + nodes.append(transpose) + + return (nodes, self._net_sources, self._parameters) + + +@_register_op(MaxPool2dOpr, AvgPool2dOpr) +class Pooling2DConverter(OperatorBaseConverter): + support_op_map = { + "AVERAGE": "AveragePool", + "AVERAGE_COUNT_EXCLUDE_PADDING": "AveragePool", + "MAX": "MaxPool", + } + + def __init__(self, opr): + super().__init__(opr) + assert ( + opr.mode.upper() in self.support_op_map + ), "Pooling op doesn't support mode {}, you can implement it in Pooling2DConverter".format( + type(opr) + ) + self.exclude_pad = opr.mode.upper() == "AVERAGE_COUNT_EXCLUDE_PADDING" + self.__opr_type__ = self.support_op_map[opr.mode.upper()] + + def _get_attrs(self): + opr = self._opr + kh, kw = expand(opr.kernel_size) + ph, pw = expand(opr.padding) + sh, sw = expand(opr.stride) + attribute = { + "kernel_shape": [kh, kw], + "pads": [ph, pw, ph, pw], + "strides": [sh, sw], + } + + if self.__opr_type__ == "AveragePool": + attribute["count_include_pad"] = 0 if self.exclude_pad else 1 + + return attribute + + +@_register_op(BatchNormalizationOpr) +class BatchnormConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + inputs = self._get_inputs(exclude_idx=[1, 2, 3, 4]) + outputs = self._get_outputs() + scale_ = opr.inp_tensors[1].np_data.squeeze() + bias_ = opr.inp_tensors[2].np_data.squeeze() + mean_ = opr.inp_tensors[3].np_data.squeeze() + var_ = opr.inp_tensors[4].np_data.squeeze() + + inputs[1] = self._opr.inp_tensors[0].name + "_scale_onnx" + inputs[2] = self._opr.inp_tensors[0].name + "_bias_onnx" + inputs[3] = self._opr.inp_tensors[0].name + "_mean_onnx" + inputs[4] = self._opr.inp_tensors[0].name + "_var_onnx" + scale = onnx.helper.make_tensor_value_info( + inputs[1], + mge2onnx_dtype_mapping[self._opr.inp_tensors[1].dtype], + scale_.shape, + ) + bias = onnx.helper.make_tensor_value_info( + inputs[2], + mge2onnx_dtype_mapping[self._opr.inp_tensors[2].dtype], + bias_.shape, + ) + mean = onnx.helper.make_tensor_value_info( + inputs[3], + mge2onnx_dtype_mapping[self._opr.inp_tensors[3].dtype], + mean_.shape, + ) + var = onnx.helper.make_tensor_value_info( + inputs[4], + mge2onnx_dtype_mapping[self._opr.inp_tensors[4].dtype], + var_.shape, + ) + self._parameters.extend( + [ + onnx.numpy_helper.from_array(scale_, inputs[1]), + onnx.numpy_helper.from_array(bias_, inputs[2]), + onnx.numpy_helper.from_array(mean_, inputs[3]), + onnx.numpy_helper.from_array(var_, inputs[4]), + ] + ) + bn = onnx.helper.make_node( + "BatchNormalization", inputs, [outputs[self._opr.output_idx]] + ) + return ([bn], self._net_sources + [scale, bias, mean, var], self._parameters) + + +@_register_op(ConcatOpr) +class ConcatConverter(OperatorBaseConverter): + __opr_type__ = "Concat" + + def __init__(self, opr): + super().__init__(opr) + if opset_version < 11: + assert ( + self._opr.axis >= 0 + ), "opset {} doesn't support negative aixs in concat opr".format( + opset_version + ) + + def _get_attrs(self): + return {"axis": self._opr.axis} + + +@_register_op(ReduceOpr) +class ReduceConverter(OperatorBaseConverter): + support_op_map = { + "MAX": "ReduceMax", + "SUM": "ReduceSum", + "SUM_SQR": "ReduceSumSquare", + } + + def __init__(self, opr): + super().__init__(opr) + assert ( + opr.mode in self.support_op_map + ), "Reduce op doesn't support mode {}, you can implement it in ReduceConverter".format( + opr.mode + ) + self.__opr_type__ = self.support_op_map[opr.mode] + + def _get_attrs(self): + if self._opr.axis < 2000000000: + return {"axes": [self._opr.axis], "keepdims": self._opr.keep_dims} + else: + return {"axes": [0], "keepdims": self._opr.keep_dims} + + def convert(self): + if self._opr.inp_tensors[0].shape == self._opr.out_tensors[0].shape: + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = onnx.helper.make_node( + self.__opr_type__, [inputs[0]], outputs, **self._get_attrs() + ) + return [nodes], self._net_sources, self._parameters + else: + inputs = self._get_inputs() + outputs = self._get_outputs() + if len(inputs) > 1: + temp_node = inputs[0] + "_reshape_in" + else: + temp_node = outputs[0] + out_nodes = [] + nodes = onnx.helper.make_node( + self.__opr_type__, [inputs[0]], [temp_node], **self._get_attrs() + ) + out_nodes.append(nodes) + if len(inputs) > 1: + shape = inputs[1] + "_shape" + shape_tensor = onnx.helper.make_tensor_value_info( + shape, + mge2onnx_dtype_mapping[np.int64], + self._opr.inp_tensors[1].shape, + ) + shape_param = onnx.numpy_helper.from_array( + self._opr.inp_tensors[1].np_data.astype(np.int64), shape + ) + self._net_sources.append(shape_tensor) + self._parameters.append(shape_param) + reshape_node = onnx.helper.make_node( + "Reshape", [temp_node, shape], outputs, + ) + out_nodes.append(reshape_node) + return out_nodes, self._net_sources, self._parameters + + +@_register_op(SqueezeOpr) +class SqueezeConverter(OperatorBaseConverter): + def convert(self): + inputs = self._get_inputs() + outputs = self._get_outputs() + remove_axis = [] + for axis in self._opr.squeeze_dims: + remove_axis.append(axis) + + if len(remove_axis) > 0: + squeeze = onnx.helper.make_node( + "Squeeze", inputs, outputs, axes=remove_axis + ) + ret = [squeeze] + else: + ret = [] + return ret, self._net_sources, self._parameters + + +@_register_op(BroadcastOpr) +class BroadcastOprConverter(OperatorBaseConverter): + def convert(self): + assert opset_version > 7, "onnx support Expand (broadcast) since opset 8" + inputs = self._get_inputs() + typecvt_node = onnx.helper.make_node( + "Cast", + [inputs[1]], + [inputs[1] + "_int64"], + to=mge2onnx_dtype_mapping[np.int64], + ) + inputs[1] = inputs[1] + "_int64" + outputs = self._get_outputs() + broadcast_node = onnx.helper.make_node("Expand", inputs, outputs) + return [typecvt_node, broadcast_node], self._net_sources, self._parameters + + +@_register_op(TypeCvtOpr) +class TypeCvtOprConverter(OperatorBaseConverter): + def convert(self): + inputs = self._get_inputs() + outputs = self._get_outputs() + target_dtype = self._opr.out_tensors[0].dtype + node = onnx.helper.make_node( + "Cast", inputs, outputs, to=mge2onnx_dtype_mapping[target_dtype] + ) + return [node], self._net_sources, self._net_sources + + +@_register_op(ConstantOpr) +class ConstantOprConverter(OperatorBaseConverter): + def convert(self): + return [], self._net_sources, self._parameters + + +@_register_op(DropoutOpr) +class DropoutOprConverter(OperatorBaseConverter): + __opr_type__ = "Dropout" + + +@_register_op(FlattenOpr) +class FlattenOprConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + assert opr.end_axis == -1, "Onnx only support end_axis = -1" + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = [] + inp_shape = list(opr.inp_tensors[0].shape) + if inp_shape[opr.start_axis] != 1: + flatten = onnx.helper.make_node( + "Flatten", inputs=inputs, outputs=outputs, axis=opr.start_axis, + ) + nodes.append(flatten) + else: + tmp_name = inputs[0] + "_tmp_flatten" + flatten = onnx.helper.make_node( + "Flatten", inputs=inputs, outputs=[tmp_name], axis=opr.start_axis, + ) + nodes.append(flatten) + squeeze = onnx.helper.make_node( + "Squeeze", [tmp_name], outputs, axes=[opr.start_axis] + ) + nodes.append(squeeze) + return nodes, self._net_sources, self._parameters + + +@_register_op(AdaptiveAvgPool2dOpr) +class AdaptiveAvgPool2dOprConverter(OperatorBaseConverter): + __opr_type__ = "AveragePool" + + def _get_attrs(self): + opr = self._opr + oh, ow = expand(opr.out_shape) + ih, iw = list(opr.inp_tensors[0].shape)[-2:] + + ph, pw = 0, 0 + sh, sw = ih // oh, iw // ow + kh, kw = ih - (oh - 1) * sh, iw - (ow - 1) * sw + + attribute = { + "kernel_shape": [kh, kw], + "pads": [ph, pw, ph, pw], + "strides": [sh, sw], + "count_include_pad": 1, + } + + return attribute + + +def gen_relu6_node(opr, inputs, outputs): + nodes = [] + if opset_version >= 11: + zero_name = opr.out_tensors[0].name + "_const_zero" + zero = onnx.helper.make_node( + "Constant", + inputs=[], + outputs=[zero_name], + value=onnx.helper.make_tensor( + name=zero_name, + data_type=onnx.TensorProto.FLOAT, # pylint: disable=no-member + dims=[], + vals=[0.0], + ), + ) + six_name = opr.out_tensors[0].name + "_const_six" + six = onnx.helper.make_node( + "Constant", + inputs=[], + outputs=[six_name], + value=onnx.helper.make_tensor( + name=six_name, + data_type=onnx.TensorProto.FLOAT, # pylint: disable=no-member + dims=[], + vals=[6.0], + ), + ) + relu6 = onnx.helper.make_node( + "Clip", inputs=[inputs[0], zero_name, six_name], outputs=outputs, + ) + nodes.extend([zero, six, relu6]) + else: + relu6 = onnx.helper.make_node( + "Clip", inputs=inputs, outputs=outputs, min=0.0, max=6.0, + ) + nodes.append(relu6) + + return nodes + + +def relu6_add3_dev6(opr, inputs, outputs): + nodes = [] + three_name = opr.out_tensors[0].name + "_const_three" + three = onnx.helper.make_node( + "Constant", + inputs=[], + outputs=[three_name], + value=onnx.helper.make_tensor( + name=three_name, + data_type=mge2onnx_dtype_mapping[opr.inp_tensors[0].dtype], + dims=[], + vals=[3.0], + ), + ) + nodes.append(three) + # x + 3 + add_three = opr.inp_tensors[0].name + "_add_three" + add_three_node = onnx.helper.make_node( + "Add", inputs=[inputs[0], three_name], outputs=[add_three] + ) + nodes.append(add_three_node) + # relu6(x+3) + relu6_out = add_three + "_relu6out" + relu6_nodes = gen_relu6_node(opr, [add_three], [relu6_out]) + nodes.extend(relu6_nodes) + # relu6(x+3)/6 + relu6_six = relu6_out + "_six" + six_name = opr.out_tensors[0].name + "_const_six" + six = onnx.helper.make_node( + "Constant", + inputs=[], + outputs=[relu6_six], + value=onnx.helper.make_tensor( + name=six_name, + data_type=mge2onnx_dtype_mapping[opr.inp_tensors[0].dtype], + dims=[], + vals=[6.0], + ), + ) + nodes.append(six) + relu6_six_dev = outputs[0] + dev6 = onnx.helper.make_node( + "Div", inputs=[relu6_out, relu6_six], outputs=[relu6_six_dev], + ) + nodes.append(dev6) + return nodes + + +@_register_op(Relu6Opr) +class Relu6OprConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = gen_relu6_node(opr, inputs, outputs) + return nodes, self._net_sources, self._parameters + + +@_register_op(HardSwishOpr) +class HardSwishOprConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = [] + relu6_six_dev = opr.inp_tensors[0].name + "_add_three_relu6out_dev6" + temp_nodes = relu6_add3_dev6(opr, inputs, [relu6_six_dev]) + nodes.extend(temp_nodes) + relu6_mul = onnx.helper.make_node( + "Mul", inputs=[inputs[0], relu6_six_dev], outputs=outputs, + ) + nodes.append(relu6_mul) + return nodes, self._net_sources, self._parameters + + +@_register_op(HardSigmoidOpr) +class HardSigmoidOprConverter(OperatorBaseConverter): + __opr_type__ = "HardSigmoid" + + def _get_attrs(self): + return { + "alpha": 1.0 / 6.0, + "beta": 0.5, + } + + def convert(self): + nodes = [ + onnx.helper.make_node( + "HardSigmoid", + self._get_inputs()[:1], + self._get_outputs(), + name=self._opr.out_tensors[0].name, + **self._get_attrs(), + ) + ] + return nodes, self._net_sources, self._parameters + + +@_register_op(RepeatOpr) +class RepeatConverter(OperatorBaseConverter): + def convert(self): + opr = self._opr + inputs = self._get_inputs() + outputs = self._get_outputs() + nodes = [] + + unsqueeze_out = inputs[0] + "_unsqueeze_out" + unsqueeze = onnx.helper.make_node( + "Unsqueeze", + inputs=[inputs[0]], + outputs=[unsqueeze_out], + axes=[opr.axis + 1], + ) + nodes.append(unsqueeze) + + repeat_shape = [1] * (opr.inp_tensors[0].ndim + 1) + repeat_shape[opr.axis + 1] *= opr.repeats + tile_repeats = unsqueeze_out + "_repeats" + tile_repeats_tensor = onnx.helper.make_tensor_value_info( + tile_repeats, + mge2onnx_dtype_mapping[np.int64], + (opr.inp_tensors[0].ndim + 1,), + ) + tile_repeats_param = onnx.numpy_helper.from_array( + np.array(repeat_shape).astype("int64"), tile_repeats + ) + self._net_sources.append(tile_repeats_tensor) + self._parameters.append(tile_repeats_param) + + repeat_name = inputs[0] + "_tile" + repeat = onnx.helper.make_node( + "Tile", inputs=[unsqueeze_out, tile_repeats], outputs=[repeat_name], + ) + nodes.append(repeat) + shape_tensor_name_after = repeat_name + "_reshape_after" + shape_tensor_after = onnx.helper.make_tensor_value_info( + shape_tensor_name_after, + mge2onnx_dtype_mapping[np.int64], + (opr.out_tensors[0].ndim,), + ) + shape_param_after = onnx.numpy_helper.from_array( + np.array(opr.out_tensors[0].shape, dtype=np.int64), shape_tensor_name_after + ) + self._net_sources.append(shape_tensor_after) + self._parameters.append(shape_param_after) + reshape_out = onnx.helper.make_node( + "Reshape", [repeat_name, shape_tensor_name_after], outputs + ) + nodes.append(reshape_out) + return nodes, self._net_sources, self._parameters diff --git a/mgeconvert/backend/ir_to_tflite/__init__.py b/mgeconvert/backend/ir_to_tflite/__init__.py new file mode 100644 index 0000000..fa9860a --- /dev/null +++ b/mgeconvert/backend/ir_to_tflite/__init__.py @@ -0,0 +1,9 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from .tflite_converter import TFLiteConverter +from .tflite_op import set_platform, set_quantization diff --git a/mgeconvert/backend/ir_to_tflite/init.sh b/mgeconvert/backend/ir_to_tflite/init.sh new file mode 100755 index 0000000..1dcdd43 --- /dev/null +++ b/mgeconvert/backend/ir_to_tflite/init.sh @@ -0,0 +1,68 @@ +#!/bin/bash -e +basepath=$(cd `dirname $0`; pwd) + + +which flatc && FLATC_VERSION="$(flatc --version)" || FLATC_VERSION="" +echo ${FLATC_VERSION} +if python3 -c "import flatbuffers">/dev/null 2>&1; then + FLAT_BUFFER_VERSION="$(python3 -m pip show flatbuffers | grep Version)" +else + FLAT_BUFFER_VERSION="" +fi +echo ${FLAT_BUFFER_VERSION} + + +if [[ "$FLATC_VERSION" != "flatc version 1.12.0" || "$FLAT_BUFFER_VERSION" != "Version: 1.12" ]]; then + sudo rm -rf /tmp/flatbuffers + git clone https://github.com/google/flatbuffers.git -b v1.12.0 /tmp/flatbuffers +fi + +if [[ "$FLATC_VERSION" != "flatc version 1.12.0" ]]; then + sudo python3 -m pip uninstall flatbuffers -y + sudo rm -rf /usr/local/bin/flatc + sudo rm -rf /usr/local/lib/libflatbuffers* + # build flatbuffers + echo "building flatbuffers..." + cd /tmp/flatbuffers + cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release -DFLATBUFFERS_BUILD_SHAREDLIB=on -DCMAKE_INSTALL_PREFIX=/usr/local + make; sudo make install +fi + + +export PATH=$PATH:/usr/local/bin +# build tflite interface from schema.fbs +echo "building tflite schema..." +cd /tmp +sudo rm -f schema.fbs +tf_version=$1 +if [ ! -n "$1" ] ;then + tf_version="r2.3" +fi +echo "Use TFLite $tf_version!" +wget https://raw.githubusercontent.com/tensorflow/tensorflow/$tf_version/tensorflow/lite/schema/schema.fbs +sudo flatc --python schema.fbs +sudo chmod 755 /tmp/tflite +cp -r /tmp/tflite $basepath + +# build pyflatbuffers +if [[ "$FLAT_BUFFER_VERSION" != "Version: 1.12" ]]; then + sudo python3 -m pip uninstall flatbuffers -y + echo "building pyflexbuffers..." + export VERSION=1.12 + cd /tmp/flatbuffers/python + sudo python3 setup.py install +fi + + +sudo python3 -m pip install pybind11==2.6.2 + +# try to find libflatbuffers.so +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib + +# using pyflexbuffers +cd $basepath/pyflexbuffers +PYBIND11_HEADER=$(python3 -c "import pybind11; print(pybind11.get_include())") +PYTHON_INCLUDE=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['include'])") +PYTHON_STDLIB=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['stdlib'])") + +g++ fbconverter.cc --std=c++14 -fPIC --shared -I${PYBIND11_HEADER} -I${PYTHON_INCLUDE} -L${PYTHON_STDLIB} -lflatbuffers -o fbconverter.so diff --git a/mgeconvert/tflite_converter/pyflexbuffers/__init__.py b/mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py similarity index 100% rename from mgeconvert/tflite_converter/pyflexbuffers/__init__.py rename to mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py diff --git a/mgeconvert/tflite_converter/pyflexbuffers/fbconverter.cc b/mgeconvert/backend/ir_to_tflite/pyflexbuffers/fbconverter.cc similarity index 100% rename from mgeconvert/tflite_converter/pyflexbuffers/fbconverter.cc rename to mgeconvert/backend/ir_to_tflite/pyflexbuffers/fbconverter.cc diff --git a/mgeconvert/tflite_converter/tflite_converter.py b/mgeconvert/backend/ir_to_tflite/tflite_converter.py similarity index 66% rename from mgeconvert/tflite_converter/tflite_converter.py rename to mgeconvert/backend/ir_to_tflite/tflite_converter.py index 7e2bd88..dac2932 100644 --- a/mgeconvert/tflite_converter/tflite_converter.py +++ b/mgeconvert/backend/ir_to_tflite/tflite_converter.py @@ -1,26 +1,23 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error import flatbuffers from tqdm import tqdm -from ..mge_context import ( - TopologyNetwork, - TransformerRule, - optimize_for_conversion, - set_platform, -) -from ..mge_context.mge_op import ( - Host2DeviceCopyOpr, +from ...converter_ir.ir_graph import IRGraph +from ...converter_ir.ir_op import ( + ConstantOpr, + LinspaceOpr, MultipleDeviceTensorHolderOpr, SharedDeviceTensorOpr, ) -from .tflite import ( # pylint: disable=import-error +from .tflite import ( Buffer, Model, Operator, @@ -29,54 +26,30 @@ SubGraph, Tensor, ) -from .tflite.CustomOptionsFormat import ( # pylint: disable=import-error - CustomOptionsFormat, +from .tflite.CustomOptionsFormat import CustomOptionsFormat +from .tflite_op import ( + MGE2TFLITE, + get_shape_param, + mge2tflite_dtype_mapping, + set_quantization, ) -from .tflite_op import MGE2TFLITE, get_shape_param, mge2tflite_dtype_mapping class TFLiteConverter: - transformer_options = [ - TransformerRule.REDUCE_AXIS_AS_INPUT, - TransformerRule.REMOVE_RESHAPE_INPUT, - TransformerRule.FUSE_FOR_RELU6, - TransformerRule.FUSE_ACTIVATION, - TransformerRule.PADDING_FOR_CONV, - TransformerRule.FUSE_FOR_CONV_BIAS, - # In TFLite Converter, ``RESHAPE_BIAS_TO_1DIM`` is required after ``FUSE_FOR_CONV_BIAS``. - TransformerRule.RESHAPE_BIAS_TO_1DIM, - # ``CONV_ADD_ZERO_BIAS`` should be after ``FUSE_FOR_CONV_BIAS`` and ``RESHAPE_BIAS_TO_1DIM``. - TransformerRule.CONV_ADD_ZERO_BIAS, - TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT, - TransformerRule.FUSE_SOFTMAX, - TransformerRule.DECONV_SHAPE_AS_INPUT, - TransformerRule.FUSE_ASTYPE, - TransformerRule.TRANSPOSE_PATTERN_AS_INPUT, - TransformerRule.FUSE_FOR_LEAKY_RELU, - TransformerRule.EXPAND_MUL_ADD3, - TransformerRule.EXPAND_ADD_SIGMOID, - TransformerRule.SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE, - ] - - def __init__(self, toponet, transformer_options=None, graph_name="graph"): - assert isinstance( - toponet, TopologyNetwork - ), "net must be instance of TopologyNetwork" - self.net = toponet + def __init__(self, net, graph_name="graph", quantizer=None): + assert isinstance(net, IRGraph), "net must be instance of IRGraph" + self.net = net self.graph_name = graph_name self._var2tensor = dict() # varnode to tensor index self._opr_type_list = [] self._buffer_list = [] self._tensor_list = [] self._operator_list = [] - self._params = {} + self.quantizer = quantizer # buffer size will automatically increase if needed self._builder = flatbuffers.Builder(1024) - - if transformer_options is not None: - self.transformer_options = transformer_options - optimize_for_conversion(self.net, self.transformer_options) + set_quantization(require_quantize=quantizer.require_quantize) def convert(self, disable_nhwc=False): # Note the 0th entry of this array must be an empty buffer (sentinel) @@ -85,17 +58,18 @@ def convert(self, disable_nhwc=False): self._buffer_list.append(buffer) def need_convert(mge_opr): - is_const = [data.np_data is not None for data in mge_opr.inp_vars] if isinstance( mge_opr, ( - Host2DeviceCopyOpr, + ConstantOpr, + LinspaceOpr, MultipleDeviceTensorHolderOpr, SharedDeviceTensorOpr, ), ): return False - return not all(is_const) or len(mge_opr.inp_vars) == 0 + is_const = [data.np_data is not None for data in mge_opr.inp_tensors] + return not all(is_const) and len(mge_opr.inp_tensors) > 0 for mge_opr in tqdm(self.net.all_oprs): last_opr = mge_opr @@ -108,44 +82,44 @@ def need_convert(mge_opr): if tfl_opr_type not in self._opr_type_list: self._opr_type_list.append(tfl_opr_type) - if hasattr(mge_opr, "type") and mge_opr.type == "ConvolutionBackwardData": - mge_opr.inp_vars = [mge_opr.inp_vars[0]] + list( - reversed(mge_opr.inp_vars[-2:]) - ) # shape, weight, input - # buffer and tensor - for var in mge_opr.inp_vars + mge_opr.out_vars: - if var in self._var2tensor: + for tensor in mge_opr.inp_tensors + mge_opr.out_tensors: + if tensor in self._var2tensor: continue - - result_shape, byte_list = get_shape_param(var, mge_opr, disable_nhwc) - var.shape = result_shape + result_shape, byte_list = get_shape_param( + tensor, mge_opr, self.quantizer, disable_nhwc=disable_nhwc + ) scale = None zero_point = 0 - if hasattr(var.dtype, "metadata"): - scale = var.dtype.metadata["mgb_dtype"]["scale"] - zero_point = var.dtype.metadata["mgb_dtype"].get("zero_point") or 0 - dtype = var.dtype - if var.name in self._params.keys(): - dtype = self._params[var.name]["dtype"] - scale = self._params[var.name]["scale"] - zero_point = self._params[var.name]["zero"] + if self.quantizer.require_quantize: + if hasattr(tensor, "scale"): + scale = tensor.scale + if hasattr(tensor, "zero_point") and tensor.zero_point is not None: + zero_point = int(tensor.zero_point) + dtype = tensor.q_dtype + from megengine.core.tensor.dtype import ( # pylint: disable=import-outside-toplevel,no-name-in-module + QuantDtypeMeta, + ) + + if isinstance(dtype, QuantDtypeMeta): + dtype = dtype.np_dtype_str + else: + dtype = tensor.dtype buffer = self.gen_buffer(byte_list) self._buffer_list.append(buffer) - tfl_tensor = self.gen_tensor( - var.name, - var.shape, + tensor.name, + result_shape, mge2tflite_dtype_mapping[dtype], len(self._buffer_list) - 1, scale=scale, - zero_point=zero_point, + zero_point=int(zero_point), ) self._tensor_list.append(tfl_tensor) - self._var2tensor[var] = len(self._tensor_list) - 1 + self._var2tensor[tensor] = len(self._tensor_list) - 1 tfl_opr = self.gen_operator( mge_opr, tfl_opr_type, tfl_options_type, tfl_options @@ -153,11 +127,11 @@ def need_convert(mge_opr): self._operator_list.append(tfl_opr) print("last op: {}".format(last_opr)) - out_var = last_opr.out_vars[0] - print("dtype: {}".format(out_var.dtype)) - if hasattr(out_var.dtype, "metadata"): - scale = out_var.dtype.metadata["mgb_dtype"]["scale"] - zero_point = out_var.dtype.metadata["mgb_dtype"].get("zero_point") or 0 + out_tensor = last_opr.out_tensors[0] + print("dtype: {}".format(out_tensor.dtype)) + if hasattr(out_tensor.dtype, "metadata"): + scale = out_tensor.dtype.metadata["mgb_dtype"]["scale"] + zero_point = out_tensor.dtype.metadata["mgb_dtype"].get("zero_point") or 0 print("scale: {}, zero point: {}".format(scale, zero_point)) return self.get_model() @@ -215,15 +189,15 @@ def gen_operator(self, opr, opr_type, options_type, options): # opcode_index opcode_index = self._opr_type_list.index(opr_type) # inputs - Operator.OperatorStartInputsVector(self._builder, len(opr.inp_vars)) - for var in reversed(opr.inp_vars): + Operator.OperatorStartInputsVector(self._builder, len(opr.inp_tensors)) + for var in reversed(opr.inp_tensors): self._builder.PrependInt32(self._var2tensor[var]) - inputs = self._builder.EndVector(len(opr.inp_vars)) + inputs = self._builder.EndVector(len(opr.inp_tensors)) # outputs - Operator.OperatorStartOutputsVector(self._builder, len(opr.out_vars)) - for var in reversed(opr.out_vars): + Operator.OperatorStartOutputsVector(self._builder, len(opr.out_tensors)) + for var in reversed(opr.out_tensors): self._builder.PrependInt32(self._var2tensor[var]) - outputs = self._builder.EndVector(len(opr.out_vars)) + outputs = self._builder.EndVector(len(opr.out_tensors)) custom_options = None builtin_options = None @@ -261,8 +235,14 @@ def get_description(self): def get_operator_codes(self): operator_codes_list = [] for opr_type in self._opr_type_list: + is_custom = not isinstance(opr_type, int) + if is_custom: + custom_code = self._builder.CreateString(opr_type.code) + opr_type = opr_type.type OperatorCode.OperatorCodeStart(self._builder) OperatorCode.OperatorCodeAddBuiltinCode(self._builder, opr_type) + if is_custom: + OperatorCode.OperatorCodeAddCustomCode(self._builder, custom_code) operator_code = OperatorCode.OperatorCodeEnd(self._builder) operator_codes_list.append(operator_code) @@ -283,16 +263,16 @@ def get_subgraphs(self): tensors = self._builder.EndVector(len(self._tensor_list)) # inputs - SubGraph.SubGraphStartInputsVector(self._builder, len(self.net.input_vars)) - for var in reversed(self.net.input_vars): + SubGraph.SubGraphStartInputsVector(self._builder, len(self.net.graph_inputs)) + for var in reversed(self.net.graph_inputs): self._builder.PrependInt32(self._var2tensor[var]) - graph_inputs = self._builder.EndVector(len(self.net.input_vars)) + graph_inputs = self._builder.EndVector(len(self.net.graph_inputs)) # outputs - SubGraph.SubGraphStartOutputsVector(self._builder, len(self.net.output_vars)) - for var in reversed(self.net.output_vars): + SubGraph.SubGraphStartOutputsVector(self._builder, len(self.net.graph_outputs)) + for var in reversed(self.net.graph_outputs): self._builder.PrependInt32(self._var2tensor[var]) - graph_outputs = self._builder.EndVector(len(self.net.output_vars)) + graph_outputs = self._builder.EndVector(len(self.net.graph_outputs)) # operators SubGraph.SubGraphStartOperatorsVector(self._builder, len(self._operator_list)) @@ -342,36 +322,3 @@ def get_model(self): model = Model.ModelEnd(self._builder) self._builder.Finish(model, "TFL3".encode()) return self._builder.Output() - - -def convert_to_tflite( - mge_fpath, output="out.tflite", *, graph_name="graph", batch_size=1, mtk=False -): - """ - Convert megengine model to TFLite, - and save the TFLite model to file `output`. - - :param mge_fpath: the file path of megengine model. - :type fpath: str - :param output: the filename used for the saved model. - :type output: str - :param graph_name: the name of the TFLite graph. - :type graph_name: str - :param batch_size: batch size of TFLite model. - :type batch_size: int - :param mtk: if this TFLite will be run on mtk. - :type mtk: bool - """ - assert isinstance(mge_fpath, str), "mge_fpath must be string" - net = TopologyNetwork(mge_fpath, prune_reshape=True) - net.batch_size = batch_size - if mtk: - # MTK devices only support batch_size 1 - net.batch_size = 1 - set_platform("mtk") - converter = TFLiteConverter(net, None, graph_name) - model = converter.convert() - - assert isinstance(output, str), "tflite_fpath must be string" - with open(output, "wb") as fout: - fout.write(model) diff --git a/mgeconvert/tflite_converter/tflite_op.py b/mgeconvert/backend/ir_to_tflite/tflite_op.py similarity index 54% rename from mgeconvert/tflite_converter/tflite_op.py rename to mgeconvert/backend/ir_to_tflite/tflite_op.py index a4f72e0..5b033f5 100644 --- a/mgeconvert/tflite_converter/tflite_op.py +++ b/mgeconvert/backend/ir_to_tflite/tflite_op.py @@ -1,45 +1,59 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error import collections from typing import List import numpy as np from numpy import dtype -from ..mge_context import ( +from ...converter_ir.ir_op import ( + AbsOpr, + AddOpr, + AvgPool2dOpr, AxisAddRemoveOpr, ConcatOpr, - ConvBiasForwardOpr, - ConvForwardBiasOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - DimshuffleOpr, - ElemwiseMultiTypeOpr, - ElemwiseOpr, + Conv2dOpr, + Deconv2dOpr, + ExpOpr, + FuseMulAdd3Opr, + GetSubTensorOpr, LeakyReluOpr, - MatrixMulOpr, - MgeOpr, + LinearOpr, + MatMulOpr, + MaxOpr, + MaxPool2dOpr, + MinOpr, + MulOpr, + OpBase, PadOpr, - PoolingForwardOpr, + PowOpr, ReduceOpr, + ReluOpr, ReshapeOpr, - ResizeForwardOpr, + ResizeOpr, + SigmoidOpr, SoftmaxOpr, SqueezeOpr, - SubtensorOpr, - Tensor, - get_platform, + SubOpr, + TransposeOpr, + TrueDivOpr, + TypeCvtOpr, ) -from .pyflexbuffers import dumps # pylint: disable=import-error -from .tflite import ( # pylint: disable=import-error +from ...converter_ir.ir_quantizer import IRQuantizer +from ...converter_ir.ir_tensor import IOHWFormat, IRTensor, NCHWFormat, OIHWFormat +from ...converter_ir.ir_transform import cal_pad_mode +from .pyflexbuffers import dumps +from .tflite import ( AbsOptions, AddOptions, + CastOptions, ConcatenationOptions, Conv2DOptions, DepthwiseConv2DOptions, @@ -49,7 +63,6 @@ LeakyReluOptions, MaximumMinimumOptions, MulOptions, - NegOptions, PadOptions, Pool2DOptions, PowOptions, @@ -63,16 +76,38 @@ TransposeConvOptions, TransposeOptions, ) -from .tflite.ActivationFunctionType import ( # pylint: disable=import-error - ActivationFunctionType, -) -from .tflite.BuiltinOperator import BuiltinOperator # pylint: disable=import-error -from .tflite.BuiltinOptions import BuiltinOptions # pylint: disable=import-error -from .tflite.Padding import Padding # pylint: disable=import-error -from .tflite.TensorType import TensorType # pylint: disable=import-error +from .tflite.ActivationFunctionType import ActivationFunctionType +from .tflite.BuiltinOperator import BuiltinOperator +from .tflite.BuiltinOptions import BuiltinOptions +from .tflite.Padding import Padding +from .tflite.TensorType import TensorType + + +class Config: + platform = "official" + require_quantize = True + + +def set_platform(platform): + assert platform in ["official", "mtk"] + Config.platform = platform + + +def set_quantization(require_quantize): + Config.require_quantize = require_quantize -def get_shape_param(tensor: Tensor, mge_opr: MgeOpr, disable_nhwc=False): +def get_platform(): + return Config.platform + + +def get_quantization(): + return Config.require_quantize + + +def get_shape_param( + tensor: IRTensor, mge_opr: OpBase, quantizer: IRQuantizer, disable_nhwc=False +): """ Return a tuple of shape and bytes(1dim) object for tflite operator, which will restore its inp/out at runtime by the shape and bytes. @@ -81,43 +116,47 @@ def get_shape_param(tensor: Tensor, mge_opr: MgeOpr, disable_nhwc=False): return tensor.shape, None shape = list(tensor.shape) - if tensor.ndim == 4 and not tensor.is_faked: + if tensor.axis_order and tensor.ndim == 4: # OC, IC, H, W to OC, H, W, IC # NCHW to NHWC # except the output of reshape if not disable_nhwc: if ( - hasattr(mge_opr, "type") - and mge_opr.type == "ConvolutionBackwardData" - and tensor.np_data is not None + isinstance(tensor.axis_order, OIHWFormat) + and mge_opr.name == "Conv2d" + and mge_opr.groups > 1 # type: ignore ): - shape = [ - tensor.shape[1], - tensor.shape[2], - tensor.shape[3], - tensor.shape[0], - ] - else: - shape = [ - tensor.shape[0], - tensor.shape[2], - tensor.shape[3], - tensor.shape[1], - ] + # Filter in DepthwiseConv is expected to be [1, H, W, O]. + shape = tensor.axis_order.shape_to_IHWO(shape) + elif isinstance(tensor.axis_order, NCHWFormat): + shape = tensor.axis_order.shape_to_NHWC(shape) + elif isinstance(tensor.axis_order, IOHWFormat): + shape = tensor.axis_order.shape_to_OHWI(shape) elif tensor.ndim > 4: assert False, "ERROR: output ndim {0} is not supported now".format(tensor.ndim) - if tensor.is_faked: - return shape, tensor.byte_list - number_list: List[np.ndarray] = [] - value = tensor.np_data + if ( + quantizer.require_quantize + and hasattr(tensor, "scale") + and tensor.np_data is not None + ): + value = quantizer.quantize(tensor) + else: + value = tensor.np_data if value is not None: - if value.ndim == 4: - if hasattr(mge_opr, "type") and mge_opr.type == "ConvolutionBackwardData": - value = value.transpose(1, 2, 3, 0) - else: - value = value.transpose(0, 2, 3, 1) + if not disable_nhwc and tensor.axis_order and value.ndim == 4: + if ( + isinstance(tensor.axis_order, OIHWFormat) + and mge_opr.name == "Conv2d" + and mge_opr.groups > 1 # type: ignore + ): + # Filter in DepthwiseConv is expected to be [1, H, W, O]. + value = tensor.axis_order.data_to_IHWO(value) + elif isinstance(tensor.axis_order, NCHWFormat): + value = tensor.axis_order.data_to_NHWC(value) + elif isinstance(tensor.axis_order, IOHWFormat): + value = tensor.axis_order.data_to_OHWI(value) number_list = value.reshape(-1) if len(number_list) > 0: @@ -134,11 +173,19 @@ def get_shape_param(tensor: Tensor, mge_opr: MgeOpr, disable_nhwc=False): np.float32: TensorType.FLOAT32, np.float16: TensorType.FLOAT16, np.int32: TensorType.INT32, + np.int16: TensorType.INT16, np.int8: TensorType.INT8, np.uint8: TensorType.UINT8, dtype("int32"): TensorType.INT32, + dtype("int16"): TensorType.INT16, dtype("uint8"): TensorType.UINT8, dtype("int8"): TensorType.INT8, + "quint8": TensorType.UINT8, + "qint32": TensorType.INT32, + "uint8": TensorType.UINT8, + "int8": TensorType.INT8, + "int16": TensorType.INT16, + "int32": TensorType.INT32, } @@ -161,169 +208,94 @@ def callback(impl): return callback -@_register_op(ElemwiseOpr, ElemwiseMultiTypeOpr) -def _elemwise(mge_opr, builder): # pylint: disable=too-many-return-statements - if isinstance(mge_opr, ElemwiseMultiTypeOpr): - # TODO: currently only support ADD + RELU - AddOptions.AddOptionsStart(builder) - AddOptions.AddOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type["RELU"] - ) - options = AddOptions.AddOptionsEnd(builder) - return BuiltinOperator.ADD, BuiltinOptions.AddOptions, options - - # return tuple of (tfl_op_type, option type, option) - if mge_opr.mode == "ABS": - AbsOptions.AbsOptionsStart(builder) - options = AbsOptions.AbsOptionsEnd(builder) - return BuiltinOperator.ABS, BuiltinOptions.AbsOptions, options - if mge_opr.mode == "NEG": - NegOptions.NegOptionsStart(builder) - options = NegOptions.NegOptionsEnd(builder) - return BuiltinOperator.NEG, BuiltinOptions.NegOptions, options - if mge_opr.mode in ("ADD", "FUSE_ADD_RELU"): - AddOptions.AddOptionsStart(builder) - AddOptions.AddOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type[mge_opr.activation] - ) - options = AddOptions.AddOptionsEnd(builder) - return BuiltinOperator.ADD, BuiltinOptions.AddOptions, options - if mge_opr.mode == "SUB": - SubOptions.SubOptionsStart(builder) - SubOptions.SubOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type[mge_opr.activation] - ) - options = SubOptions.SubOptionsEnd(builder) - return BuiltinOperator.SUB, BuiltinOptions.SubOptions, options - if mge_opr.mode == "MUL": - MulOptions.MulOptionsStart(builder) - MulOptions.MulOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type[mge_opr.activation] - ) - options = MulOptions.MulOptionsEnd(builder) - return BuiltinOperator.MUL, BuiltinOptions.MulOptions, options - if mge_opr.mode in ("DIV", "TRUE_DIV"): - DivOptions.DivOptionsStart(builder) - DivOptions.DivOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type[mge_opr.activation] - ) - options = DivOptions.DivOptionsEnd(builder) - return BuiltinOperator.DIV, BuiltinOptions.DivOptions, options - if mge_opr.mode == "POW": - PowOptions.PowOptionsStart(builder) - options = PowOptions.PowOptionsEnd(builder) - return BuiltinOperator.POW, BuiltinOptions.PowOptions, options - if mge_opr.mode == "EXP": - ExpOptions.ExpOptionsStart(builder) - options = ExpOptions.ExpOptionsEnd(builder) - return BuiltinOperator.EXP, BuiltinOptions.ExpOptions, options - if mge_opr.mode == "MAX": - MaximumMinimumOptions.MaximumMinimumOptionsStart(builder) - options = MaximumMinimumOptions.MaximumMinimumOptionsEnd(builder) - return BuiltinOperator.MAXIMUM, BuiltinOptions.MaximumMinimumOptions, options - if mge_opr.mode == "MIN": - MaximumMinimumOptions.MaximumMinimumOptionsStart(builder) - options = MaximumMinimumOptions.MaximumMinimumOptionsEnd(builder) - return BuiltinOperator.MINIMUM, BuiltinOptions.MaximumMinimumOptions, options - if mge_opr.mode == "SIGMOID": - return BuiltinOperator.LOGISTIC, None, None - return None, None, None +@_register_op(AddOpr, FuseMulAdd3Opr) +def _add(mge_opr, builder): # pylint: disable=too-many-return-statements + AddOptions.AddOptionsStart(builder) + AddOptions.AddOptionsAddFusedActivationFunction( + builder, mge2tflite_activation_type[mge_opr.activation] + ) + options = AddOptions.AddOptionsEnd(builder) + return BuiltinOperator.ADD, BuiltinOptions.AddOptions, options -@_register_op(ReduceOpr) -def _reduce(mge_opr, builder): - ReducerOptions.ReducerOptionsStart(builder) - ReducerOptions.ReducerOptionsAddKeepDims(builder, True) - options = ReducerOptions.ReducerOptionsEnd(builder) +@_register_op(SubOpr) +def _sub(mge_opr, builder): + SubOptions.SubOptionsStart(builder) + SubOptions.SubOptionsAddFusedActivationFunction( + builder, mge2tflite_activation_type[mge_opr.activation] + ) + options = SubOptions.SubOptionsEnd(builder) + return BuiltinOperator.SUB, BuiltinOptions.SubOptions, options - reduce_mode_map = { - "SUM": BuiltinOperator.SUM, - "MEAN": BuiltinOperator.MEAN, - "MAX": BuiltinOperator.REDUCE_MAX, - "MIN": BuiltinOperator.REDUCE_MIN, - } - return reduce_mode_map[mge_opr.mode], BuiltinOptions.ReducerOptions, options +@_register_op(AbsOpr) +def _abs(_, builder): + AbsOptions.AbsOptionsStart(builder) + options = AbsOptions.AbsOptionsEnd(builder) + return BuiltinOperator.ABS, BuiltinOptions.AbsOptions, options -@_register_op(ReshapeOpr, AxisAddRemoveOpr) -def _reshape(mge_opr, builder): - ReshapeOptions.ReshapeOptionsStartNewShapeVector(builder, len(mge_opr.output_shape)) - for i in reversed(list(mge_opr.output_shape)): - builder.PrependInt32(i) - new_shape = builder.EndVector(len(mge_opr.output_shape)) - ReshapeOptions.ReshapeOptionsStart(builder) - ReshapeOptions.ReshapeOptionsAddNewShape(builder, new_shape) - options = ReshapeOptions.ReshapeOptionsEnd(builder) - return BuiltinOperator.RESHAPE, BuiltinOptions.ReshapeOptions, options +@_register_op(SigmoidOpr) +def _sigmoid(*_): # pylint: disable=too-many-return-statements + return BuiltinOperator.LOGISTIC, None, None -@_register_op(ConcatOpr) -def _concat(mge_opr, builder): - ConcatenationOptions.ConcatenationOptionsStart(builder) - ConcatenationOptions.ConcatenationOptionsAddFusedActivationFunction( + +@_register_op(ReluOpr) +def _relu(*_): + return BuiltinOperator.RELU, None, None + + +@_register_op(MulOpr) +def _mul(mge_opr, builder): # pylint: disable=too-many-return-statements + MulOptions.MulOptionsStart(builder) + MulOptions.MulOptionsAddFusedActivationFunction( builder, mge2tflite_activation_type[mge_opr.activation] ) - axis = mge_opr.axis - if mge_opr.inp_vars[0].ndim == 4: - # map NCHW to NHWC - if mge_opr.axis == 1: - axis = 3 - elif mge_opr.axis == 2: - axis = 1 - elif mge_opr.axis == 3: - axis = 2 - ConcatenationOptions.ConcatenationOptionsAddAxis(builder, axis) - options = ConcatenationOptions.ConcatenationOptionsEnd(builder) - return BuiltinOperator.CONCATENATION, BuiltinOptions.ConcatenationOptions, options + options = MulOptions.MulOptionsEnd(builder) + return BuiltinOperator.MUL, BuiltinOptions.MulOptions, options -@_register_op(PoolingForwardOpr) -def _pooling(mge_opr, builder): - Pool2DOptions.Pool2DOptionsStart(builder) - Pool2DOptions.Pool2DOptionsAddPadding(builder, Padding.VALID) - shape = mge_opr.inp_vars[0].shape - if get_platform() == "mtk" and shape[2] == mge_opr.kh and shape[3] == mge_opr.kw: - # MTK global pooling - print( - "\nWARNING: the stride of global pooling " - "would be changed to 1 to adapted the bug of MTK" - ) - Pool2DOptions.Pool2DOptionsAddStrideH(builder, 1) - Pool2DOptions.Pool2DOptionsAddStrideW(builder, 1) - else: - Pool2DOptions.Pool2DOptionsAddStrideH(builder, mge_opr.sh) - Pool2DOptions.Pool2DOptionsAddStrideW(builder, mge_opr.sw) - Pool2DOptions.Pool2DOptionsAddFilterHeight(builder, mge_opr.kh) - Pool2DOptions.Pool2DOptionsAddFilterWidth(builder, mge_opr.kw) - Pool2DOptions.Pool2DOptionsAddFusedActivationFunction( +@_register_op(TrueDivOpr) +def _div(mge_opr, builder): + DivOptions.DivOptionsStart(builder) + DivOptions.DivOptionsAddFusedActivationFunction( builder, mge2tflite_activation_type[mge_opr.activation] ) - options = Pool2DOptions.Pool2DOptionsEnd(builder) + options = DivOptions.DivOptionsEnd(builder) + return BuiltinOperator.DIV, BuiltinOptions.DivOptions, options - tfl_opr_type = BuiltinOperator.AVERAGE_POOL_2D - if mge_opr.mode == "MAX": - tfl_opr_type = BuiltinOperator.MAX_POOL_2D - return tfl_opr_type, BuiltinOptions.Pool2DOptions, options + +def _padding_mode_conv(mge_opr): + if cal_pad_mode(mge_opr) == "VALID": + return Padding.VALID + else: + return Padding.SAME -@_register_op(ConvolutionForwardOpr, ConvBiasForwardOpr, ConvForwardBiasOpr) +@_register_op(Conv2dOpr) def _conv2d(mge_opr, builder): - if mge_opr.group > 1: + if mge_opr.groups > 1: DepthwiseConv2DOptions.DepthwiseConv2DOptionsStart(builder) - DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddPadding(builder, Padding.VALID) - DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddStrideH(builder, mge_opr.sh) - DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddStrideW(builder, mge_opr.sw) + DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddPadding( + builder, _padding_mode_conv(mge_opr) + ) + DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddStrideH( + builder, mge_opr.stride[0] + ) + DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddStrideW( + builder, mge_opr.stride[1] + ) DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddDepthMultiplier( - builder, mge_opr.inp_vars[1].shape[0] + builder, mge_opr.inp_tensors[1].shape[0] ) DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddFusedActivationFunction( builder, mge2tflite_activation_type[mge_opr.activation] ) DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddDilationHFactor( - builder, mge_opr.dilation_h + builder, mge_opr.dilation[0] ) DepthwiseConv2DOptions.DepthwiseConv2DOptionsAddDilationWFactor( - builder, mge_opr.dilation_w + builder, mge_opr.dilation[1] ) options = DepthwiseConv2DOptions.DepthwiseConv2DOptionsEnd(builder) return ( @@ -333,11 +305,11 @@ def _conv2d(mge_opr, builder): ) Conv2DOptions.Conv2DOptionsStart(builder) - Conv2DOptions.Conv2DOptionsAddPadding(builder, Padding.VALID) - Conv2DOptions.Conv2DOptionsAddStrideH(builder, mge_opr.sh) - Conv2DOptions.Conv2DOptionsAddStrideW(builder, mge_opr.sw) - Conv2DOptions.Conv2DOptionsAddDilationHFactor(builder, mge_opr.dilation_h) - Conv2DOptions.Conv2DOptionsAddDilationWFactor(builder, mge_opr.dilation_w) + Conv2DOptions.Conv2DOptionsAddPadding(builder, _padding_mode_conv(mge_opr)) + Conv2DOptions.Conv2DOptionsAddStrideH(builder, mge_opr.stride[0]) + Conv2DOptions.Conv2DOptionsAddStrideW(builder, mge_opr.stride[1]) + Conv2DOptions.Conv2DOptionsAddDilationHFactor(builder, mge_opr.dilation[0]) + Conv2DOptions.Conv2DOptionsAddDilationWFactor(builder, mge_opr.dilation[1]) Conv2DOptions.Conv2DOptionsAddFusedActivationFunction( builder, mge2tflite_activation_type[mge_opr.activation] ) @@ -345,49 +317,23 @@ def _conv2d(mge_opr, builder): return BuiltinOperator.CONV_2D, BuiltinOptions.Conv2DOptions, options -@_register_op(ResizeForwardOpr) -def _resize(_, builder): - ResizeBilinearOptions.ResizeBilinearOptionsStart(builder) - ResizeBilinearOptions.ResizeBilinearOptionsAddAlignCorners(builder, False) - options = ResizeBilinearOptions.ResizeBilinearOptionsEnd(builder) - return ( - BuiltinOperator.RESIZE_BILINEAR, - BuiltinOptions.ResizeBilinearOptions, - options, - ) - - -@_register_op(MatrixMulOpr) -def _matrix_mul(mge_opr, builder): - FullyConnectedOptions.FullyConnectedOptionsStart(builder) - # mge quantized model should not have bias for tflite conversion - FullyConnectedOptions.FullyConnectedOptionsAddFusedActivationFunction( - builder, mge2tflite_activation_type[mge_opr.activation] - ) - options = FullyConnectedOptions.FullyConnectedOptionsEnd(builder) - return ( - BuiltinOperator.FULLY_CONNECTED, - BuiltinOptions.FullyConnectedOptions, - options, - ) - - -@_register_op(SoftmaxOpr) -def _softmax(mge_opr, builder): - SoftmaxOptions.SoftmaxOptionsStart(builder) - SoftmaxOptions.SoftmaxOptionsAddBeta(builder, mge_opr.beta) - options = SoftmaxOptions.SoftmaxOptionsEnd(builder) - return BuiltinOperator.SOFTMAX, BuiltinOptions.SoftmaxOptions, options +@_register_op(PadOpr) +def _pad(_, builder): + PadOptions.PadOptionsStart(builder) + options = PadOptions.PadOptionsEnd(builder) + return BuiltinOperator.PAD, BuiltinOptions.PadOptions, options def _padding_mode_transpose_conv(mge_opr): if ( - mge_opr.out_vars[0].shape[2] == mge_opr.inp_vars[1].shape[2] * mge_opr.sh - and mge_opr.out_vars[0].shape[3] == mge_opr.inp_vars[1].shape[3] * mge_opr.sw + mge_opr.out_tensors[0].shape[3] + == mge_opr.inp_tensors[2].shape[3] * mge_opr.stride[0] + and mge_opr.out_tensors[0].shape[2] + == mge_opr.inp_tensors[2].shape[2] * mge_opr.stride[1] ): # padding mode == SAME return Padding.SAME - elif mge_opr.ph == 0 and mge_opr.pw == 0: + elif mge_opr.padding[0] == 0 and mge_opr.padding[1] == 0: # padding mode == VALID return Padding.VALID else: @@ -395,22 +341,22 @@ def _padding_mode_transpose_conv(mge_opr): return None -@_register_op(ConvolutionBackwardDataOpr) +@_register_op(Deconv2dOpr) def _deconv(mge_opr, builder): if get_platform() == "mtk": - CustomOperator = collections.namedtuple("CustomOperator", ["code"]) + CustomOperator = collections.namedtuple("CustomOperator", ["type", "code"]) CustomOptions = collections.namedtuple("CustomOptions", ["code"]) options = dict() options["PaddingType"] = _padding_mode_transpose_conv(mge_opr) - options["stride_height"] = mge_opr.sh - options["stride_width"] = mge_opr.sw + options["stride_height"] = mge_opr.stride[0] + options["stride_width"] = mge_opr.stride[1] options["depth_multiplier"] = 1 - options["dilation_height_factor"] = mge_opr.dilation_h - options["dilation_width_factor"] = mge_opr.dilation_w + options["dilation_height_factor"] = mge_opr.dilation[0] + options["dilation_width_factor"] = mge_opr.dilation[1] options["activation"] = mge2tflite_activation_type[mge_opr.activation] return ( - CustomOperator("MTK_TRANSPOSE_CONV"), + CustomOperator(type=BuiltinOperator.CUSTOM, code="MTK_TRANSPOSE_CONV"), CustomOptions("MTK_TRANSPOSE_CONV"), dumps(options), ) @@ -419,41 +365,156 @@ def _deconv(mge_opr, builder): TransposeConvOptions.TransposeConvOptionsAddPadding( builder, _padding_mode_transpose_conv(mge_opr) ) - TransposeConvOptions.TransposeConvOptionsAddStrideH(builder, mge_opr.sh) - TransposeConvOptions.TransposeConvOptionsAddStrideW(builder, mge_opr.sw) + TransposeConvOptions.TransposeConvOptionsAddStrideH(builder, mge_opr.stride[0]) + TransposeConvOptions.TransposeConvOptionsAddStrideW(builder, mge_opr.stride[1]) options = TransposeConvOptions.TransposeConvOptionsEnd(builder) return BuiltinOperator.TRANSPOSE_CONV, BuiltinOptions.TransposeConvOptions, options -@_register_op(PadOpr) -def _pad(_, builder): - PadOptions.PadOptionsStart(builder) - options = PadOptions.PadOptionsEnd(builder) - return BuiltinOperator.PAD, BuiltinOptions.PadOptions, options +@_register_op(ConcatOpr) +def _concat(mge_opr, builder): + ConcatenationOptions.ConcatenationOptionsStart(builder) + ConcatenationOptions.ConcatenationOptionsAddFusedActivationFunction( + builder, mge2tflite_activation_type[mge_opr.activation] + ) + axis = mge_opr.axis + if mge_opr.inp_tensors[0].ndim == 4: + # map NCHW to NHWC + if mge_opr.axis == 1: + axis = 3 + elif mge_opr.axis == 2: + axis = 1 + elif mge_opr.axis == 3: + axis = 2 + ConcatenationOptions.ConcatenationOptionsAddAxis(builder, axis) + options = ConcatenationOptions.ConcatenationOptionsEnd(builder) + return BuiltinOperator.CONCATENATION, BuiltinOptions.ConcatenationOptions, options -@_register_op(DimshuffleOpr) -def _dimshuffle(_, builder): - TransposeOptions.TransposeOptionsStart(builder) - options = TransposeOptions.TransposeOptionsEnd(builder) - return BuiltinOperator.TRANSPOSE, BuiltinOptions.TransposeOptions, options +@_register_op(PowOpr) +def _pow(_, builder): + PowOptions.PowOptionsStart(builder) + options = PowOptions.PowOptionsEnd(builder) + return BuiltinOperator.POW, BuiltinOptions.PowOptions, options -@_register_op(LeakyReluOpr) -def _leaky_relu(mge_opr, builder): - LeakyReluOptions.LeakyReluOptionsStart(builder) - LeakyReluOptions.LeakyReluOptionsAddAlpha(builder, mge_opr.negative_slope[0]) - options = LeakyReluOptions.LeakyReluOptionsEnd(builder) - return BuiltinOperator.LEAKY_RELU, BuiltinOptions.LeakyReluOptions, options +@_register_op(ExpOpr) +def _exp(_, builder): + ExpOptions.ExpOptionsStart(builder) + options = ExpOptions.ExpOptionsEnd(builder) + return BuiltinOperator.EXP, BuiltinOptions.ExpOptions, options + + +@_register_op(MaxOpr) +def _max(_, builder): + MaximumMinimumOptions.MaximumMinimumOptionsStart(builder) + options = MaximumMinimumOptions.MaximumMinimumOptionsEnd(builder) + return BuiltinOperator.MAXIMUM, BuiltinOptions.MaximumMinimumOptions, options + + +@_register_op(MinOpr) +def _min(_, builder): + MaximumMinimumOptions.MaximumMinimumOptionsStart(builder) + options = MaximumMinimumOptions.MaximumMinimumOptionsEnd(builder) + return BuiltinOperator.MINIMUM, BuiltinOptions.MaximumMinimumOptions, options -@_register_op(SubtensorOpr) -def _subtensor(_, builder): +@_register_op(ReshapeOpr, AxisAddRemoveOpr) +def _reshape(mge_opr, builder): + ReshapeOptions.ReshapeOptionsStartNewShapeVector(builder, len(mge_opr.out_shape)) + for i in reversed(list(mge_opr.out_shape)): + builder.PrependInt32(i) + new_shape = builder.EndVector(len(mge_opr.out_shape)) + ReshapeOptions.ReshapeOptionsStart(builder) + ReshapeOptions.ReshapeOptionsAddNewShape(builder, new_shape) + options = ReshapeOptions.ReshapeOptionsEnd(builder) + return BuiltinOperator.RESHAPE, BuiltinOptions.ReshapeOptions, options + + +@_register_op(ResizeOpr) +def _resize(mge_opr, builder): + assert mge_opr.mode == "bilinear", "Resize mode should be BILINEAR." + ResizeBilinearOptions.ResizeBilinearOptionsStart(builder) + ResizeBilinearOptions.ResizeBilinearOptionsAddAlignCorners(builder, False) + ResizeBilinearOptions.ResizeBilinearOptionsAddHalfPixelCenters(builder, True) + options = ResizeBilinearOptions.ResizeBilinearOptionsEnd(builder) + return ( + BuiltinOperator.RESIZE_BILINEAR, + BuiltinOptions.ResizeBilinearOptions, + options, + ) + + +@_register_op(ReduceOpr) +def _reduce(mge_opr, builder): + ReducerOptions.ReducerOptionsStart(builder) + ReducerOptions.ReducerOptionsAddKeepDims(builder, False) + options = ReducerOptions.ReducerOptionsEnd(builder) + + reduce_mode_map = { + "SUM": BuiltinOperator.SUM, + "MEAN": BuiltinOperator.MEAN, + "MAX": BuiltinOperator.REDUCE_MAX, + "MIN": BuiltinOperator.REDUCE_MIN, + } + return reduce_mode_map[mge_opr.mode], BuiltinOptions.ReducerOptions, options + + +@_register_op(GetSubTensorOpr) +def _getsubtensor(_, builder): StridedSliceOptions.StridedSliceOptionsStart(builder) options = StridedSliceOptions.StridedSliceOptionsEnd(builder) return BuiltinOperator.STRIDED_SLICE, BuiltinOptions.StridedSliceOptions, options +@_register_op(MaxPool2dOpr, AvgPool2dOpr) +def _pooling(mge_opr, builder): + Pool2DOptions.Pool2DOptionsStart(builder) + Pool2DOptions.Pool2DOptionsAddPadding(builder, Padding.VALID) + shape = mge_opr.inp_tensors[0].shape + if ( + get_platform() == "mtk" + and shape[2] == mge_opr.kernel_size[0] + and shape[3] == mge_opr.kernel_size[1] + ): + # MTK global pooling + print( + "\nWARNING: the stride of global pooling " + "would be changed to 1 to adapted the bug of MTK" + ) + Pool2DOptions.Pool2DOptionsAddStrideH(builder, 1) + Pool2DOptions.Pool2DOptionsAddStrideW(builder, 1) + else: + Pool2DOptions.Pool2DOptionsAddStrideH(builder, mge_opr.stride[0]) + Pool2DOptions.Pool2DOptionsAddStrideW(builder, mge_opr.stride[1]) + Pool2DOptions.Pool2DOptionsAddFilterHeight(builder, mge_opr.kernel_size[0]) + Pool2DOptions.Pool2DOptionsAddFilterWidth(builder, mge_opr.kernel_size[1]) + Pool2DOptions.Pool2DOptionsAddFusedActivationFunction( + builder, mge2tflite_activation_type[mge_opr.activation] + ) + options = Pool2DOptions.Pool2DOptionsEnd(builder) + + tfl_opr_type = BuiltinOperator.AVERAGE_POOL_2D + if mge_opr.name == "MaxPool2d": + tfl_opr_type = BuiltinOperator.MAX_POOL_2D + return tfl_opr_type, BuiltinOptions.Pool2DOptions, options + + +@_register_op(MatMulOpr, LinearOpr) +def _matrix_mul(mge_opr, builder): + FullyConnectedOptions.FullyConnectedOptionsStart(builder) + # mge quantized model should not have bias for tflite conversion + FullyConnectedOptions.FullyConnectedOptionsAddFusedActivationFunction( + builder, mge2tflite_activation_type[mge_opr.activation] + ) + options = FullyConnectedOptions.FullyConnectedOptionsEnd(builder) + return ( + BuiltinOperator.FULLY_CONNECTED, + BuiltinOptions.FullyConnectedOptions, + options, + ) + + @_register_op(SqueezeOpr) def _squeeze(mge_opr, builder): SqueezeOptions.SqueezeOptionsStartSqueezeDimsVector( @@ -466,3 +527,42 @@ def _squeeze(mge_opr, builder): SqueezeOptions.SqueezeOptionsAddSqueezeDims(builder, squeeze_dims) options = SqueezeOptions.SqueezeOptionsEnd(builder) return BuiltinOperator.SQUEEZE, BuiltinOptions.SqueezeOptions, options + + +@_register_op(TransposeOpr) +def _transpose(_, builder): + TransposeOptions.TransposeOptionsStart(builder) + options = TransposeOptions.TransposeOptionsEnd(builder) + return BuiltinOperator.TRANSPOSE, BuiltinOptions.TransposeOptions, options + + +@_register_op(SoftmaxOpr) +def _softmax(mge_opr, builder): + SoftmaxOptions.SoftmaxOptionsStart(builder) + SoftmaxOptions.SoftmaxOptionsAddBeta(builder, mge_opr.beta) + options = SoftmaxOptions.SoftmaxOptionsEnd(builder) + return BuiltinOperator.SOFTMAX, BuiltinOptions.SoftmaxOptions, options + + +@_register_op(LeakyReluOpr) +def _leaky_relu(mge_opr, builder): + LeakyReluOptions.LeakyReluOptionsStart(builder) + LeakyReluOptions.LeakyReluOptionsAddAlpha(builder, mge_opr.negative_slope) + options = LeakyReluOptions.LeakyReluOptionsEnd(builder) + return BuiltinOperator.LEAKY_RELU, BuiltinOptions.LeakyReluOptions, options + + +@_register_op(TypeCvtOpr) +def _typecvt(mge_opr, builder): + if get_quantization(): + target_type = mge_opr.inp_tensors[0].q_dtype + else: + target_type = mge_opr.inp_tensors[0].dtype + + CastOptions.CastOptionsStart(builder) + CastOptions.CastOptionsAddInDataType(builder, mge2tflite_dtype_mapping[target_type]) + CastOptions.CastOptionsAddOutDataType( + builder, mge2tflite_dtype_mapping[target_type] + ) + options = CastOptions.CastOptionsEnd(builder) + return BuiltinOperator.CAST, BuiltinOptions.CastOptions, options diff --git a/mgeconvert/caffe_converter/caffe_converter.py b/mgeconvert/caffe_converter/caffe_converter.py deleted file mode 100644 index fa472a3..0000000 --- a/mgeconvert/caffe_converter/caffe_converter.py +++ /dev/null @@ -1,151 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from google.protobuf import text_format # type: ignore[attr-defined] -from tqdm import tqdm - -from ..mge_context import TopologyNetwork, TransformerRule, optimize_for_conversion -from ..mge_context.mge_utils import get_symvar_value -from .caffe_op import MGE2CAFFE -from .caffe_pb import caffe_pb2 as cp # pylint: disable=import-error - - -class CaffeConverter: - transformer_options = [ - TransformerRule.FUSE_FOR_LEAKY_RELU, - TransformerRule.FUSE_FOR_CONV_BIAS, - TransformerRule.FUSE_FOR_DECONV_BIAS, - TransformerRule.FUSE_FOR_FULLY_CONNECTED, - ] - - def __init__(self, toponet, transform_options=None, use_empty_blobs=False): - self.net = toponet - self.var2blob_map = {} - self.layers = [] - self._names = set() - self._count = 0 - self.use_empty_blobs = use_empty_blobs - - if transform_options is not None: - self.transformer_options = transform_options - optimize_for_conversion(self.net, self.transformer_options) - - def dump(self, proto_file, caffe_file=None): - CaffeNet = cp.NetParameter(layer=self.layers) - if caffe_file is not None: - with open(caffe_file, "wb") as f: - f.write(CaffeNet.SerializeToString()) - - for layer in CaffeNet.layer: - layer.ClearField("blobs") - - with open(proto_file, "w") as f: - f.write(text_format.MessageToString(CaffeNet)) - - @property - def gen_name(self): - self._count = self._count + 1 - while "_caffe_{0}".format(self._count) in self._names: - self._count = self._count + 1 - return "_caffe_{0}".format(self._count) - - def get_blob_name(self, varNode): - if varNode not in self.var2blob_map: - raise KeyError("can not find VarNode {}".format(varNode)) - return self.var2blob_map[varNode] - - def set_blob_name(self, varNode, name=None): - assert varNode not in self.var2blob_map, "{} already be set".format(varNode) - if name is not None: - assert isinstance(name, str) - self.var2blob_map[varNode] = name - else: - self.var2blob_map[varNode] = self.gen_name - self._names.add(self.var2blob_map[varNode]) - return self.var2blob_map[varNode] - - def reset_blob_name(self, varNode, name=None): - assert varNode in self.var2blob_map, "{} should be set".format(varNode) - if name is not None: - assert isinstance(name, str) - self.var2blob_map[varNode] = name - else: - self.var2blob_map[varNode] = self.gen_name - self._names.add(self.var2blob_map[varNode]) - return self.var2blob_map[varNode] - - def gen_blob_proto(self, data): - if self.use_empty_blobs: - return cp.BlobProto() - if isinstance(data, (int, float)): - return cp.BlobProto(data=[data]) - else: - return cp.BlobProto( - data=data.reshape(-1), shape=cp.BlobShape(dim=data.shape) - ) - - def add_layer(self, layer): - if isinstance(layer, list): - for x in layer: - self.layers.append(x) - else: - self.layers.append(layer) - - def convert(self): - unsupported_oprs = [] - for opr in self.net.all_oprs: - if not isinstance(opr, tuple(MGE2CAFFE.keys())): - unsupported_oprs.append(opr) - continue - unsupported_oprs = set(map(type, unsupported_oprs)) - assert not unsupported_oprs, "Operators {} are not supported yet".format( - unsupported_oprs - ) - - def need_convert(opr): - is_const = [data.np_data is not None for data in opr.inp_vars] - return not all(is_const) or len(opr.inp_vars) == 0 - - all_oprs = list(self.net.all_oprs) - - for index in range(len(all_oprs) - 1, -1, -1): - if all_oprs[index].skip: - del all_oprs[index] - - for opr in tqdm(all_oprs): - if not need_convert(opr): - for tensor in opr.out_vars: - if tensor.np_data is None: - tensor.np_data = get_symvar_value(tensor._var) - continue - MGE2CAFFE[type(opr)](opr, self) - - -def convert_to_caffe( - mge_fpath, prototxt="out.prototxt", caffemodel="out.caffemodel", outspec=None -): - """ - Convert megengine model to Caffe, - and save caffe model to `prototxt` and `caffemodel`. - - :param mge_fpath: the file path of megengine model. - :type mge_fpath: str - :param prototxt: the filename used for saved model definition. - :type prototxt: str - :param caffemodel: the filename used for saved model weights. - :type caffemodel: str - """ - - assert isinstance(mge_fpath, str), "mge_fpath must be string" - net = TopologyNetwork(mge_fpath, outspec=outspec) - converter = CaffeConverter(net) - converter.convert() - assert isinstance(prototxt, str) and isinstance( - caffemodel, str - ), "'prototxt' and 'caffemodel' must be string" - converter.dump(prototxt, caffemodel) diff --git a/mgeconvert/cambricon_converter/.gitignore b/mgeconvert/cambricon_converter/.gitignore deleted file mode 100644 index 5d78563..0000000 --- a/mgeconvert/cambricon_converter/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -build/* -lib/cnlib/cambriconLib.py -lib/cnlib/_cambriconLib.so diff --git a/mgeconvert/cambricon_converter/CMakeLists.txt b/mgeconvert/cambricon_converter/CMakeLists.txt deleted file mode 100644 index e85e50d..0000000 --- a/mgeconvert/cambricon_converter/CMakeLists.txt +++ /dev/null @@ -1,42 +0,0 @@ -cmake_minimum_required(VERSION 3.9.0) -project(CambriconSDK) - -set(CMAKE_CXX_STANDARD 14) -set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_CXX_EXTENSIONS OFF) -set(CMAKE_POSITION_INDEPENDENT_CODE ON) - -include(cmake/cnrt.cmake) -include(cmake/cndev.cmake) -include(cmake/cnml.cmake) -list(APPEND CAMBRICON_LIBS libcnrt libcndev libcnml) -set(CAMBRICON_LIBS "${CAMBRICON_LIBS}") -set(CAMBRICON_INCLUDE_DIR "${CNML_INCLUDE_DIR}") - -find_package(PythonLibs ${PYTHON_VERSION_STRING} REQUIRED) - -find_package(SWIG REQUIRED) -set(SDK_OUTPUT_LIB_NAME cambriconLib) -set(SWIG_SDK_SRC swig/cambricon.i) -set(CMAKE_SWIG_FLAGS -Wall -threads -py3 -modern -DSWIGWORDSIZE64 -I${CAMBRICON_INCLUDE_DIR}) - -include(UseSWIG) -set_property(SOURCE ${SWIG_SDK_SRC} PROPERTY CPLUSPLUS ON) -swig_add_library(${SDK_OUTPUT_LIB_NAME} LANGUAGE python SOURCES ${SWIG_SDK_SRC}) - -set_target_properties(_${SDK_OUTPUT_LIB_NAME} PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${_sdk_name}) -target_include_directories(_${SDK_OUTPUT_LIB_NAME} PRIVATE ${PYTHON_INCLUDE_DIRS} ${CAMBRICON_INCLUDE_DIR} ${CMAKE_CURRENT_BINARY_DIR}) -target_link_libraries(_${SDK_OUTPUT_LIB_NAME} ${PYTHON_LIBRARIES} ${CAMBRICON_LIBS}) - -if (TARGET _${SDK_OUTPUT_LIB_NAME}) - add_custom_target( - develop - COMMAND ${CMAKE_COMMAND} -E create_symlink - ${CMAKE_CURRENT_BINARY_DIR}/$ - ${CMAKE_CURRENT_SOURCE_DIR}/lib/cnlib/$ - COMMAND ${CMAKE_COMMAND} -E create_symlink - ${CMAKE_CURRENT_BINARY_DIR}/${SDK_OUTPUT_LIB_NAME}.py - ${CMAKE_CURRENT_SOURCE_DIR}/lib/cnlib/${SDK_OUTPUT_LIB_NAME}.py - DEPENDS _${SDK_OUTPUT_LIB_NAME} - ) -endif() diff --git a/mgeconvert/cambricon_converter/README.md b/mgeconvert/cambricon_converter/README.md deleted file mode 100644 index e698acd..0000000 --- a/mgeconvert/cambricon_converter/README.md +++ /dev/null @@ -1,29 +0,0 @@ -# 寒武纪转换器 - -## 模型要求 - -寒武纪模型在MLU270上推理时要求含乘法的算子采取量化(`QInt8`)计算，包括`Conv`、`MatMul`、`Lrn`等，其他算子采用浮点数(`Float32`或`Float16`)计算。因此转换前的MegEngine模型也要具备这个特征才能保证转换后模型的精度不下降。 - -对应算子的输入和输出类型要求： - -- 量化算子：src(qint8) + filter(qint8) + bias(qint32) -> dst(qint32) -- 其他算子：src(fp32) -> dst(fp32) - -参考[测试用例](test/quantization_utils.py)中的`QuantizationLinearOpr`和`QuantizationConvBnOpr`来完成此类模型的搭建。 - - -## 添加未支持算子 - -模型转换分成两个阶段，增加算子也遵循这个规则。 - -- 解析MegEngine模型，可参考其他算子的实现 - 1. 在上级目录 `mge_context/mge_op.py` 中增加相应的算子 -- 搭建寒武纪模型 - 1. 在本目录 `swig/cambricon.i` 中封装寒武纪SDK提供的算子接口 - 2. 在本目录 `lib/operators.py` 中实现该算子类 - 3. 在本目录 `converter.py` 中增加相应的转换函数 - 4. 在根目录 `test/utils` 中添加测试 - -## 对分经验 - -如果转换成功但对分失败，可在文件`converter.py`的`convert`函数中设置`end_op`进行二分法对分，快速定位是错误算子。 \ No newline at end of file diff --git a/mgeconvert/cambricon_converter/cambricon_converter.py b/mgeconvert/cambricon_converter/cambricon_converter.py deleted file mode 100644 index 08ce7a1..0000000 --- a/mgeconvert/cambricon_converter/cambricon_converter.py +++ /dev/null @@ -1,235 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from typing import List - -import numpy as np - -from ..mge_context import ( - ConcatOpr, - ConvBiasForwardOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - MatrixMulOpr, - TopologyNetwork, - TransformerRule, - optimize_for_conversion, -) -from ..mge_context.mge_utils import get_dtype_name, get_logger -from .cambricon_op import MGE2CN -from .lib import cnq -from .lib import operators as cnop -from .lib.model import Model -from .lib.tensor import DATA_TYPE, TENSOR_TYPE, Tensor - -logger = get_logger(__name__) - - -map_dtype = { - "float16": DATA_TYPE.FLOAT16, - "float32": DATA_TYPE.FLOAT32, -} - - -class CambriconConverter: - transformer_options: List[TransformerRule] = [] - - def __init__( - self, - net, - transformer_options=None, - batch_size=4, - core_number=1, - data_type="float32", - use_nhwc=False, - ): - if use_nhwc: - Tensor.NCHW2NHWC = True - Tensor._default_dtype = map_dtype[data_type] - self.batch_size = batch_size - self.core_number = core_number - self.data_type = map_dtype[data_type] - self.mge_net = net - if transformer_options is not None: - self.transformer_options = transformer_options - optimize_for_conversion(self.mge_net, self.transformer_options) - self.var_map = {} - self.cn_inputs = [] - self.cn_outputs = [] - self.cn_oprs = [] - self.fusion = None - - self.check_model() - - def check_model(self): - logger.info("check model...") - unsupported_oprs = set() - quantization_error_oprs = set() - for opr in self.mge_net.all_oprs: - if not isinstance(opr, tuple(MGE2CN.keys())): - unsupported_oprs.add(type(opr)) - if isinstance( - opr, - ( - ConvBiasForwardOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - MatrixMulOpr, - ), - ): - if ( - get_dtype_name(opr.inp_vars[0]) != "QuantizedS8" - or get_dtype_name(opr.inp_vars[1]) != "QuantizedS8" - or ( - len(opr.inp_vars) > 2 - and get_dtype_name(opr.inp_vars[2]) != "QuantizedS32" - ) - or get_dtype_name(opr.out_vars[0]) != "QuantizedS32" - ): - quantization_error_oprs.add(type(opr)) - - if unsupported_oprs: - logger.error("Operators %s are not supported yet.", unsupported_oprs) - if quantization_error_oprs: - logger.error( - "Operators %s should be quantized, " - "check the function test_linear in test/test_cambricon for inspiration", - quantization_error_oprs, - ) - assert not unsupported_oprs and not quantization_error_oprs - - def get_cn_tensor(self, var): - if var not in self.var_map: - raise KeyError("can not find var {}".format(var.__dict__)) - return self.var_map[var] - - def set_cn_tensor(self, opr, var): - ttype = TENSOR_TYPE.CONST - dtype = self.data_type - shp = var.shape - data = var.np_data - - if var.qbit == "QuantizedS8": - ttype = TENSOR_TYPE.FILTER - dtype = DATA_TYPE.INT8 - if var.qbit == "QuantizedS32": - data = data.astype(np.float32) * var.scale - - if len(shp) == 1: # conv bias - shp = (1, shp[0], 1, 1) - data = data.reshape(shp) - if len(shp) == 5: # group conv's filter - shp = (shp[0] * shp[1], shp[2], shp[3], shp[4]) - data = data.reshape(shp) - # if var.qbit is None and shp[0] != 1: - if isinstance(opr, ConcatOpr): - shp = (self.batch_size,) + shp[1:] - data = np.broadcast_to(data[:1], shp) - return Tensor( - shape=shp, - ttype=ttype, - dtype=dtype, - data=data, - scale=var.scale, - name=opr.name, - ) - - def get_cn_inputs_and_outputs(self, opr): - cn_inps = [] - for var in opr.inp_vars: - if var not in self.var_map: - self.var_map[var] = self.set_cn_tensor(opr, var) - cn_inps.append(self.var_map[var]) - cn_oups = [] - for var in opr.out_vars: - shp = (self.batch_size,) + var.shape[1:] - self.var_map[var] = Tensor(shape=shp, name=var.name, scale=var.scale) - cn_oups.append(self.var_map[var]) - return cn_inps, cn_oups - - def add_oprs(self, *cnoprs): - self.cn_oprs.extend(cnoprs) - - def convert(self, end_op=None): - for opr in self.mge_net.all_oprs: - # Prune operators which calculate parameters. - pruning = True - for var in opr.out_vars: - pruning = False if var.np_data is None else pruning - if pruning: - continue - MGE2CN[type(opr)](opr, self) - if opr.name == end_op: - end_op = opr - break - assert not isinstance(end_op, str), ( - 'This operator does not exist: "%s"' % end_op - ) - self.cn_inputs = self.cn_oprs[0].inputs[0] - if end_op is None: - self.cn_outputs = self.cn_oprs[-1].outputs[0] - else: - self.cn_outputs = self.var_map[end_op.out_vars[0]] - - def fuse(self): - self.fusion = cnop.Fusion("fusion", self.cn_inputs, self.cn_outputs) - for cnopr in self.cn_oprs: - self.fusion.fuse_op(cnopr) - self.fusion.set_fusion_io() - self.fusion.set_core_num(self.core_number) - self.fusion.compile() - - def forward(self, feed_input: "np.ndarray"): - self.cn_inputs.cpudata = feed_input - self.cn_inputs.h2d() - self.fusion.forward(cnq) - return self.cn_outputs.cpudata - - def dump(self, fname): - model = Model(self.fusion.name) - model.add_fusionop(self.fusion) - model.dump(fname) - - -def convert_to_cambricon( - mge_fpath, filename, batch_size, core_number, data_type, use_nhwc -): - """ - Convert megengine model to cambricon model. - - :param mge_fpath: the file path of megengine model. - :type mge_fpath: str - :param filename: cambricon model file name. - :type filename: str - :param batch_size: batch_size of the output cambricon model. - :type batch_size: int - :param core_number: core_number of the output cambricon model. - :type core_number: int - :param data_type: data_type of the output cambricon model, which should be - "float32" or "float16". - :type data_type: str - :param use_nhwc: whether to use nhwc layout. - :type use_nhwc: bool - """ - assert isinstance(mge_fpath, str), "mge_fpath must be string" - net = TopologyNetwork(mge_fpath) - logger.info("init converter...") - converter = CambriconConverter( - net, - batch_size=batch_size, - core_number=core_number, - data_type=data_type, - use_nhwc=use_nhwc, - ) - logger.info("convert operators to cambricon...") - converter.convert() - logger.info("%d operators converted...", len(converter.cn_oprs)) - converter.fuse() - logger.info("fusing...") - converter.dump(filename) - logger.info("ok, dump model to %s", filename) diff --git a/mgeconvert/cambricon_converter/cambricon_op.py b/mgeconvert/cambricon_converter/cambricon_op.py deleted file mode 100644 index a6c89b1..0000000 --- a/mgeconvert/cambricon_converter/cambricon_op.py +++ /dev/null @@ -1,340 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import numpy as np - -from ..mge_context import ( - AxisAddRemoveOpr, - BatchedMatrixMulOpr, - BatchNormForwardOpr, - BroadcastOpr, - ConcatOpr, - ConvBiasForwardOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - DimshuffleOpr, - ElemwiseMultiTypeOpr, - ElemwiseOpr, - GetVarShapeOpr, - Host2DeviceCopyOpr, - IdentityOpr, - MarkNoBroadcastElemwiseOpr, - MatrixMulOpr, - MultipleDeviceTensorHolderOpr, - PoolingForwardOpr, - ReduceOpr, - ReshapeOpr, - SharedDeviceTensorOpr, - SubtensorOpr, - TypeCvtOpr, -) -from .lib import operators as cnop -from .lib.tensor import TENSOR_TYPE, Tensor - -MGE2CN = {} - - -def _register_op(*oprs): - def callback(impl): - for opr in oprs: - MGE2CN[opr] = impl - return impl - - return callback - - -@_register_op(MultipleDeviceTensorHolderOpr) -def _multiple_device_tensor_holder(*_): - pass - - -@_register_op(SharedDeviceTensorOpr) -def _shared_device_tensor(*_): - pass - - -@_register_op(GetVarShapeOpr) -def _get_var_shape(*_): - pass - - -@_register_op(Host2DeviceCopyOpr) -def _host2device_copy(opr, context): - context.var_map[opr.out_vars[0]] = Tensor((context.batch_size,) + opr.shape[1:]) - - -@_register_op(TypeCvtOpr) -def _typecvt(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - context.var_map[opr.out_vars[0]] = inps[0] - if oups[0].scale is not None: - inps[0].scale = oups[0].scale - - -@_register_op(ReshapeOpr) -def _reshape(opr, context): - opr.inp_vars = opr.inp_vars[:1] - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.Reshape(opr.name, inps, oups) - context.add_oprs(cnopr) - - -@_register_op(BroadcastOpr) -def _broad_cast(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.Broadcast(opr.name, inps, oups) - context.add_oprs(cnopr) - - -@_register_op(DimshuffleOpr) -def _dimshuffle(opr, context): - if opr.ndim == 1: - inps, _ = context.get_cn_inputs_and_outputs(opr) - context.var_map[opr.out_vars[0]] = inps[0] - else: - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.Dimshuffle(opr.name, inps, oups, opr.pattern) - context.add_oprs(cnopr) - - -@_register_op(ReduceOpr) -def _reduce(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - if opr.mode == "SUM_SQR": - # Reduce(SUM_SQR) -> Elemwise(MUL) + Reduce(SUM) - cnt_square = Tensor(shape=inps[0].shape) - cn_square = cnop.Square(opr.name + "_SQUARE", inps, cnt_square) - cn_sum = cnop.Reduce(opr.name + "_SUM", cnt_square, oups, opr.axis, "SUM") - context.add_oprs(cn_square, cn_sum) - else: - cnopr = cnop.Reduce(opr.name, inps, oups, opr.axis, opr.mode) - context.add_oprs(cnopr) - - -@_register_op(ConcatOpr) -def _concat(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - assert inps[1].shape[0] == inps[1].shape[0] - cnopr = cnop.Concat(opr.name, inps, oups, opr.axis) - context.add_oprs(cnopr) - - -@_register_op(AxisAddRemoveOpr) -def _axis_add_remove(opr, context): - inps, _ = context.get_cn_inputs_and_outputs(opr) - context.var_map[opr.out_vars[0]] = inps[0] - - -@_register_op(MarkNoBroadcastElemwiseOpr) -def _mark_no_broadcast_elemwise(opr, context): - inps, _ = context.get_cn_inputs_and_outputs(opr) - context.var_map[opr.out_vars[0]] = inps[0] - - -@_register_op(SubtensorOpr) -def _subtensor(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - inp_shape = inps[0].shape - slc = [[0, inp_shape[i], 1] for i in range(4)] - for i in range(len(opr.axis)): - slc[opr.axis[i]] = [opr.begin_param[i], opr.end_param[i], opr.step_param[i]] - cnopr = cnop.Slice(opr.name, inps, oups, slc) - context.add_oprs(cnopr) - - -@_register_op(IdentityOpr) -def _identity(opr, context): - inps, _ = context.get_cn_inputs_and_outputs(opr) - context.var_map[opr.out_vars[0]] = inps[0] - - -@_register_op(ElemwiseOpr) -def _elemwise(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - if opr.mode == "FUSE_ADD_RELU": - # FUSE_ADD_RELU -> ADD + RELU - cnt_add = Tensor(shape=oups[0].shape) - cn_add = cnop.Elemwise(opr.name + "_ADD", inps, cnt_add, "ADD") - cn_relu = cnop.Active(opr.name + "_RELU", cnt_add, oups, "RELU") - context.add_oprs(cn_add, cn_relu) - return - if opr.mode == "TRUE_DIV": - # CNML does not support broadcast div. - if inps[1].cpudata is not None: - # TRUE_DIV -> BAISC_DIV + MUL - cnt_basic_div = Tensor( - shape=inps[1].shape, ttype=TENSOR_TYPE.CONST, data=1 / inps[1].cpudata - ) - cn_mul = cnop.Elemwise( - opr.name + "_MUL", [inps[0], cnt_basic_div], oups, "CYCLE_MUL" - ) - context.add_oprs(cn_mul) - return - else: - # real_div: support (n, c, h, w) / (1, 1, 1, 1) or (n, c, h, w) / (n, c, h, w) - cnopr = cnop.Elemwise(opr.name, inps, oups, "TRUE_DIV") - context.add_oprs(cnopr) - return - if opr.mode in ("NONE", "SIGMOID", "RELU", "TANH"): - cnopr = cnop.Active(opr.name, inps, oups, opr.mode) - elif opr.mode == "POW": - if float(inps[1].cpudata) == 2: - cnopr = cnop.Square(opr.name, inps[0], oups) - elif float(inps[1].cpudata) == 0.5: - cnopr = cnop.Sqrt(opr.name, inps[0], oups) - else: - raise NotImplementedError("power op in cambricon cannot work correctly") - elif opr.mode in ("ADD", "MUL"): - if inps[0].cpudata is None and inps[1].cpudata is None: - cnopr = cnop.Elemwise(opr.name, inps, oups, opr.mode) - else: - if inps[0].cpudata is None: - cnopr = cnop.Elemwise(opr.name, inps, oups, "CYCLE_" + opr.mode) - else: - cnopr = cnop.Elemwise( - opr.name, [inps[1], inps[0]], oups, "CYCLE_" + opr.mode - ) - elif opr.mode == "NEGATE": - inps.insert(0, Tensor((1, 1, 1, 1), TENSOR_TYPE.CONST, data=0)) - cnopr = cnop.Elemwise(opr.name, inps, oups, opr.mode) - else: - cnopr = cnop.Elemwise(opr.name, inps, oups, opr.mode) - context.add_oprs(cnopr) - - -@_register_op(ElemwiseMultiTypeOpr) -def _elemwise_multitype(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - # QADD -> ADD + RELU - cnt_add = Tensor(shape=oups[0].shape) - cn_add = cnop.Elemwise(opr.name + "Add", inps, cnt_add, "ADD") - cn_relu = cnop.Active(opr.name + "Relu", cnt_add, oups, "RELU") - context.add_oprs(cn_add, cn_relu) - - -@_register_op(MatrixMulOpr) -def _matrix_mul(opr, context): - # cnml does not support transposeB. - if not opr.transposeB: - var = opr.inp_vars[1] - var.shape = (var.shape[1], var.shape[0]) - var.np_data = np.transpose(var.np_data) - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.MatMul(opr.name, inps[0], oups) - cnopr.param_dict["weight"] = inps[1] - context.add_oprs(cnopr) - - -@_register_op(BatchedMatrixMulOpr) -def _batched_matrix_mul(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.BatchMatMul(opr.name, inps, oups) - context.add_oprs(cnopr) - - -@_register_op(BatchNormForwardOpr) -def _batch_norm(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.BatchNorm(opr.name, inps[0], oups[2]) - cnopr.param_dict["running_mean"] = Tensor( - shape=opr.mean.shape, ttype=TENSOR_TYPE.CONST, data=opr.mean - ) - cnopr.param_dict["running_var"] = Tensor( - shape=opr.var.shape, ttype=TENSOR_TYPE.CONST, data=opr.var - ) - context.add_oprs(cnopr) - - -@_register_op(ConvolutionForwardOpr) -def _convolution(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.Conv( - opr.name, - inps[0], - oups, - opr.sh, - opr.sw, - opr.dilation_h, - opr.dilation_w, - opr.ph, - opr.pw, - groups=opr.group, - ) - cnopr.param_dict["W"] = inps[1] - b_shp = (1, inps[1].shape[0], 1, 1) - cnopr.param_dict["B"] = Tensor( - shape=b_shp, ttype=TENSOR_TYPE.CONST, data=np.zeros(b_shp) - ) - context.add_oprs(cnopr) - - -@_register_op(ConvBiasForwardOpr) -def _conv_bias(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - # ConvBias -> Convolution + Active - cnt_convolution = Tensor(shape=oups[0].shape) - # Convolution - cn_convolution = cnop.Conv( - opr.name + "Conv", - inps[0], - cnt_convolution, - opr.sh, - opr.sw, - opr.dilation_h, - opr.dilation_w, - opr.ph, - opr.pw, - groups=opr.group, - ) - cn_convolution.param_dict["W"] = inps[1] - cn_convolution.param_dict["B"] = inps[2] - # Active - cn_active = cnop.Active(opr.name + "Act", cnt_convolution, oups, opr.activation) - context.add_oprs(cn_convolution, cn_active) - - -@_register_op(ConvolutionBackwardDataOpr) -def _deconv(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cn_deconv = cnop.Deconv( - opr.name, - inps[1], - oups, - opr.sh, - opr.sw, - opr.dilation_h, - opr.dilation_w, - opr.ph, - opr.pw, - ) - cn_deconv.param_dict["W"] = inps[0] - b_shp = (1, inps[0].shape[0], 1, 1) - cn_deconv.param_dict["B"] = Tensor( - shape=b_shp, ttype=TENSOR_TYPE.CONST, data=np.zeros(b_shp) - ) - context.add_oprs(cn_deconv) - - -@_register_op(PoolingForwardOpr) -def _pooling(opr, context): - inps, oups = context.get_cn_inputs_and_outputs(opr) - cnopr = cnop.Pool( - opr.name, - inps, - oups, - opr.kh, - opr.kw, - opr.sh, - opr.sw, - opr.ph, - opr.pw, - 1, - 1, - "MAX" if opr.mode == "MAX" else "AVG", - ) - context.add_oprs(cnopr) diff --git a/mgeconvert/cambricon_converter/cmake/cndev.cmake b/mgeconvert/cambricon_converter/cmake/cndev.cmake deleted file mode 100644 index 7f9a037..0000000 --- a/mgeconvert/cambricon_converter/cmake/cndev.cmake +++ /dev/null @@ -1,48 +0,0 @@ -if($ENV{LIBRARY_PATH}) - string(REPLACE ":" ";" SYSTEM_LIBRARY_PATHS $ENV{LIBRARY_PATH}) -endif() - -find_library(CNDEV_LIBRARY - NAMES libcndev.so - PATHS $ENV{LD_LIBRARY_PATH} "$ENV{NEUWARE_HOME}/lib64" ${CMAKE_INSTALL_PREFIX} - HINTS ${SYSTEM_LIBRARY_PATHS} - PATH_SUFFIXES lib lib64 - DOC "CNDEV library." ) - -if(CNDEV_LIBRARY STREQUAL "CNDEV_LIBRARY-NOTFOUND") - message(FATAL_ERROR "Can not find CNDEV Library") -endif() - -get_filename_component(__found_cndev_root "${CNDEV_LIBRARY}/../include" REALPATH) -find_path(CNDEV_INCLUDE_DIR - NAMES cndev.h - HINTS "$ENV{NEUWARE_HOME}/include" ${__found_cndev_root} - PATH_SUFFIXES include - DOC "Path to CNDEV include directory." ) - -if(CNDEV_INCLUDE_DIR STREQUAL "CNDEV_INCLUDE_DIR-NOTFOUND") - message(FATAL_ERROR "Can not find CNDEV Library") -endif() - -file(STRINGS "${CNDEV_INCLUDE_DIR}/cndev.h" CNDEV_1 REGEX "^#define CNDEV_VERSION_1 [0-9]+.*$") -file(STRINGS "${CNDEV_INCLUDE_DIR}/cndev.h" CNDEV_2 REGEX "^#define CNDEV_VERSION_2 [0-9]+.*$") -file(STRINGS "${CNDEV_INCLUDE_DIR}/cndev.h" CNDEV_3 REGEX "^#define CNDEV_VERSION_3 [0-9]+.*$") -file(STRINGS "${CNDEV_INCLUDE_DIR}/cndev.h" CNDEV_4 REGEX "^#define CNDEV_VERSION_4 [0-9]+.*$") -file(STRINGS "${CNDEV_INCLUDE_DIR}/cndev.h" CNDEV_5 REGEX "^#define CNDEV_VERSION_5 [0-9]+.*$") - -string(REGEX REPLACE "^#define CNDEV_VERSION_1 ([0-9]+).*$" "\\1" CNDEV_VERSION_1 "${CNDEV_1}") -string(REGEX REPLACE "^#define CNDEV_VERSION_2 ([0-9]+).*$" "\\1" CNDEV_VERSION_2 "${CNDEV_2}") -string(REGEX REPLACE "^#define CNDEV_VERSION_3 ([0-9]+).*$" "\\1" CNDEV_VERSION_3 "${CNDEV_3}") -string(REGEX REPLACE "^#define CNDEV_VERSION_4 ([0-9]+).*$" "\\1" CNDEV_VERSION_4 "${CNDEV_4}") -string(REGEX REPLACE "^#define CNDEV_VERSION_5 ([0-9]+).*$" "\\1" CNDEV_VERSION_5 "${CNDEV_5}") -set(CNDEV_VERSION_STRING "${CNDEV_VERSION_1}.${CNDEV_VERSION_2}.${CNDEV_VERSION_3}.${CNDEV_VERSION_4}.${CNDEV_VERSION_5}") - -add_library(libcndev SHARED IMPORTED) - -set_target_properties(libcndev PROPERTIES - IMPORTED_LOCATION ${CNDEV_LIBRARY} - INTERFACE_INCLUDE_DIRECTORIES ${CNDEV_INCLUDE_DIR} -) - -message("-- Found CNDEV: ${__found_cndev_root} (found version: ${CNDEV_VERSION_STRING})") - diff --git a/mgeconvert/cambricon_converter/cmake/cnml.cmake b/mgeconvert/cambricon_converter/cmake/cnml.cmake deleted file mode 100644 index 60649ce..0000000 --- a/mgeconvert/cambricon_converter/cmake/cnml.cmake +++ /dev/null @@ -1,44 +0,0 @@ -if($ENV{LIBRARY_PATH}) - string(REPLACE ":" ";" SYSTEM_LIBRARY_PATHS $ENV{LIBRARY_PATH}) -endif() - -find_library(CNML_LIBRARY - NAMES libcnml.so - PATHS $ENV{LD_LIBRARY_PATH} "$ENV{NEUWARE_HOME}/lib64" ${CMAKE_INSTALL_PREFIX} - HINTS ${SYSTEM_LIBRARY_PATHS} - PATH_SUFFIXES lib lib64 - DOC "CNML library." ) - -if(CNML_LIBRARY STREQUAL "CNML_LIBRARY-NOTFOUND") - message(FATAL_ERROR "Can not find CNML Library") -endif() - -get_filename_component(__found_cnml_root "${CNML_LIBRARY}/../include" REALPATH) -find_path(CNML_INCLUDE_DIR - NAMES cnml.h - HINTS "$ENV{NEUWARE_HOME}/include" ${__found_cnml_root} - PATH_SUFFIXES include - DOC "Path to CNML include directory." ) - -if(CNML_INCLUDE_DIR STREQUAL "CNML_INCLUDE_DIR-NOTFOUND") - message(FATAL_ERROR "Can not find CNML Library") -endif() - -file(STRINGS "${CNML_INCLUDE_DIR}/cnml.h" CNML_MAJOR REGEX "^#define CNML_MAJOR_VERSION [0-9]+.*$") -file(STRINGS "${CNML_INCLUDE_DIR}/cnml.h" CNML_MINOR REGEX "^#define CNML_MINOR_VERSION [0-9]+.*$") -file(STRINGS "${CNML_INCLUDE_DIR}/cnml.h" CNML_PATCH REGEX "^#define CNML_PATCH_VERSION [0-9]+.*$") - -string(REGEX REPLACE "^#define CNML_MAJOR_VERSION ([0-9]+).*$" "\\1" CNML_VERSION_MAJOR "${CNML_MAJOR}") -string(REGEX REPLACE "^#define CNML_MINOR_VERSION ([0-9]+).*$" "\\1" CNML_VERSION_MINOR "${CNML_MINOR}") -string(REGEX REPLACE "^#define CNML_PATCH_VERSION ([0-9]+).*$" "\\1" CNML_VERSION_PATCH "${CNML_PATCH}") -set(CNML_VERSION_STRING "${CNML_VERSION_MAJOR}.${CNML_VERSION_MINOR}.${CNML_VERSION_PATCH}") - -add_library(libcnml SHARED IMPORTED) - -set_target_properties(libcnml PROPERTIES - IMPORTED_LOCATION ${CNML_LIBRARY} - INTERFACE_INCLUDE_DIRECTORIES ${CNML_INCLUDE_DIR} -) - -message("-- Found CNML: ${__found_cnml_root} (found version: ${CNML_VERSION_STRING})") - diff --git a/mgeconvert/cambricon_converter/cmake/cnrt.cmake b/mgeconvert/cambricon_converter/cmake/cnrt.cmake deleted file mode 100644 index 363c070..0000000 --- a/mgeconvert/cambricon_converter/cmake/cnrt.cmake +++ /dev/null @@ -1,44 +0,0 @@ -if($ENV{LIBRARY_PATH}) - string(REPLACE ":" ";" SYSTEM_LIBRARY_PATHS $ENV{LIBRARY_PATH}) -endif() - -find_library(CNRT_LIBRARY - NAMES libcnrt.so - PATHS $ENV{LD_LIBRARY_PATH} "$ENV{NEUWARE_HOME}/lib64" ${CMAKE_INSTALL_PREFIX} - HINTS ${SYSTEM_LIBRARY_PATHS} - PATH_SUFFIXES lib lib64 - DOC "CNRT library." ) - -if(CNRT_LIBRARY STREQUAL "CNRT_LIBRARY-NOTFOUND") - message(FATAL_ERROR "Can not find CNRT Library") -endif() - -get_filename_component(__found_cnrt_root "${CNRT_LIBRARY}/../include" REALPATH) -find_path(CNRT_INCLUDE_DIR - NAMES cnrt.h - HINTS "$ENV{NEUWARE_HOME}/include" ${__found_cnrt_root} - PATH_SUFFIXES include - DOC "Path to CNRT include directory." ) - -if(CNRT_INCLUDE_DIR STREQUAL "CNRT_INCLUDE_DIR-NOTFOUND") - message(FATAL_ERROR "Can not find CNRT Library") -endif() - -file(STRINGS "${CNRT_INCLUDE_DIR}/cnrt.h" CNRT_MAJOR REGEX "^#define CNRT_MAJOR_VERSION [0-9]+.*$") -file(STRINGS "${CNRT_INCLUDE_DIR}/cnrt.h" CNRT_MINOR REGEX "^#define CNRT_MINOR_VERSION [0-9]+.*$") -file(STRINGS "${CNRT_INCLUDE_DIR}/cnrt.h" CNRT_PATCH REGEX "^#define CNRT_PATCH_VERSION [0-9]+.*$") - -string(REGEX REPLACE "^#define CNRT_MAJOR_VERSION ([0-9]+).*$" "\\1" CNRT_VERSION_MAJOR "${CNRT_MAJOR}") -string(REGEX REPLACE "^#define CNRT_MINOR_VERSION ([0-9]+).*$" "\\1" CNRT_VERSION_MINOR "${CNRT_MINOR}") -string(REGEX REPLACE "^#define CNRT_PATCH_VERSION ([0-9]+).*$" "\\1" CNRT_VERSION_PATCH "${CNRT_PATCH}") -set(CNRT_VERSION_STRING "${CNRT_VERSION_MAJOR}.${CNRT_VERSION_MINOR}.${CNRT_VERSION_PATCH}") - -add_library(libcnrt SHARED IMPORTED) - -set_target_properties(libcnrt PROPERTIES - IMPORTED_LOCATION ${CNRT_LIBRARY} - INTERFACE_INCLUDE_DIRECTORIES ${CNRT_INCLUDE_DIR} -) - -message("-- Found CNRT: ${__found_cnrt_root} (found version: ${CNRT_VERSION_STRING})") - diff --git a/mgeconvert/cambricon_converter/init.sh b/mgeconvert/cambricon_converter/init.sh deleted file mode 100755 index 08b45f1..0000000 --- a/mgeconvert/cambricon_converter/init.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash -e -hash wget || (echo "please install wget package" && exit -1) - -cd $(dirname $0) - -wget https://raw.githubusercontent.com/numpy/numpy/master/tools/swig/numpy.i -P swig/ - -mkdir build && cd build -cmake .. && make -mv _cambriconLib.so ../lib/cnlib/ -mv cambriconLib.py ../lib/cnlib/ diff --git a/mgeconvert/cambricon_converter/lib/model.py b/mgeconvert/cambricon_converter/lib/model.py deleted file mode 100644 index e9e573b..0000000 --- a/mgeconvert/cambricon_converter/lib/model.py +++ /dev/null @@ -1,23 +0,0 @@ -from .cnlib import cambriconLib as cnlib - - -class Model: - __name = None - __cnml_model = None - - def __init__(self, name): - self.__name = name - self.__cnml_model = cnlib.cnModel(self.__name) - - def dump(self, fname): - cnlib.cnmlSaveModel(self.__cnml_model, fname) - - @property - def name(self): - return self.__name - - def add_fusionop(self, fusion): - if fusion.compiled is False: - fusion.compile() - - cnlib.cnmlAddFusionOpToModel(self.__cnml_model, fusion.op, fusion.name) diff --git a/mgeconvert/cambricon_converter/lib/operators.py b/mgeconvert/cambricon_converter/lib/operators.py deleted file mode 100644 index b047361..0000000 --- a/mgeconvert/cambricon_converter/lib/operators.py +++ /dev/null @@ -1,1265 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from typing import Sequence - -import numpy as np - -from .cnlib import cambriconLib as cnlib -from .tensor import Tensor - - -class OperatorBase: - - __name = None - __inp_tensor: Sequence = [] - __out_tensor: Sequence = [] - _compiled = None - _param = None - _op = None - _op_type = None - - def __init__(self, name, optype, inputs=None, outputs=None): - self.name = name - self.inputs = inputs - self.outputs = outputs - self._compiled = False - self.param_dict = {} - self._op_type = optype - - for x in self.inputs: - x.ensure_mlu() - - @property - def name(self): - return self.__name - - @name.setter - def name(self, n): - assert isinstance(n, str), "Invalid 'name' type{}".format(type(n)) - self.__name = n - - @property - def inputs(self): - return self.__inp_tensor - - @inputs.setter - def inputs(self, inp): - if not isinstance(inp, Sequence): - inp = [inp] - self.__inp_tensor = inp - - @property - def outputs(self): - return self.__out_tensor - - @outputs.setter - def outputs(self, out): - if not isinstance(out, Sequence): - out = [out] - self.__out_tensor = out - - @property - def type(self): - return self._op_type - - @property - def compiled(self): - return self._compiled - - @property - def op(self): - return self._op - - def _ensure_param(self): - if self._param is None: - self.make_param() - - def _ensure_made(self): - if self._op is None: - self.make() - - def _ensure_compiled(self): - if not self.compiled: - self.compile() - - def make_param(self): - """Create param. - Oprs should override this function.""" - - def make(self): - """Create opr. - Oprs should not override this function.""" - - self._ensure_param() - if self._op is None: - self.make_device() - - def make_device(self): - """Create opr, please make it clear that host opr is not need to be created. - Oprs should override this function.""" - - raise NotImplementedError("Function was not implemented.") - - def compile(self): - if self._op is None: - self.make() - cnlib.cnmlCompileBaseOp_V2(self._op) - self._compiled = True - - def forward(self, cnqueue): - """Run opr on device. - Opr should not override this function.""" - - self._ensure_compiled() - self.forward_trait(cnqueue) - cnlib.cnrtSyncQueue(cnqueue) - for t in self.outputs: - t.d2h() - - def forward_trait(self, cnqueue): - """Call CNML computation API. - Opr should override this function.""" - - raise NotImplementedError("Function was not implemented.") - - def destroy_param(self): - """Destroy param. - Opr should override this function.""" - - def destroy(self): - """Destroy opr. - Opr should not override this function.""" - - if not self._op is None: - cnlib.cnDestroyBaseOp(self._op) - self._op = None - if not self._param is None: - self.destroy_param() - self._param = None - for t in self.param_dict.values(): - if isinstance(t, Tensor): - t.destroy() - self.param_dict.clear() - - -class DeviceMemcpy(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "DeviceMemcpy", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnDevMemcpyOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeDeviceMemcpyOpForward_V4( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Floor(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Floor", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnFloorOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeFloorOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Active(OperatorBase): - __active_mode = None - - __active_map = { - "NONE": cnlib.CNML_ACTIVE_NONE, - "IDENTITY": cnlib.CNML_ACTIVE_NONE, - "SIGMOID": cnlib.CNML_ACTIVE_SIGMOID, - "RELU": cnlib.CNML_ACTIVE_RELU, - "TANH": cnlib.CNML_ACTIVE_TANH, - "RELU1": cnlib.CNML_ACTIVE_RELU1, - "RELU6": cnlib.CNML_ACTIVE_RELU6, - "HARD_SIGMOID": cnlib.CNML_ACTIVE_HARD_SIGMOID, - } - - def __init__(self, name, inputs, outputs, active_mode): - super().__init__(name, "Active", inputs, outputs) - self.__active_mode = self.__active_map[active_mode] - - def make_device(self): - self._op = cnlib.cnActiveOp( - self.__active_mode, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeActiveOpForward_V4( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Elemwise(OperatorBase): - __elemwise_type = None - - __create_map = { - "ABS": cnlib.cnAbsOp, - "EXP": cnlib.cnExpOp, - "LOG": cnlib.cnLogOp, - "POW": cnlib.cnPowOp, - "ADD": cnlib.cnAddOp, - "SUB": cnlib.cnSubOp, - "MUL": cnlib.cnMultOp, - "TRUE_DIV": cnlib.cnDivOp, - "NEGATE": cnlib.cnSubOp, - "CYCLE_ADD": cnlib.cnCycleAddOp, - "CYCLE_MUL": cnlib.cnCycleMultOp, - } - - __compute_map = { - "ABS": cnlib.cnmlComputeAbsOpForward_V4, - "EXP": cnlib.cnmlComputeExpOpForward_V4, - "LOG": cnlib.cnmlComputeLogOpForward_V4, - "POW": cnlib.cnmlComputePowerOpForward_V4, - "ADD": cnlib.cnmlComputeBroadcastAddOpForward_V4, - "SUB": cnlib.cnmlComputeBroadcastSubOpForward_V4, - "MUL": cnlib.cnmlComputeBroadcastMultOpForward_V4, - "TRUE_DIV": cnlib.cnmlComputeRealDivOpForward_V4, - "NEGATE": cnlib.cnmlComputeBroadcastSubOpForward_V4, - "CYCLE_ADD": cnlib.cnmlComputeCycleAddOpForward_V4, - "CYCLE_MUL": cnlib.cnmlComputeCycleMultOpForward_V4, - } - - def __init__(self, name, inputs, outputs, elemwise_type, **kwargs): - super().__init__(name, "Elemwise", inputs, outputs) - self.__elemwise_type = elemwise_type - self._kwargs = kwargs - - def make_device(self): - if self.__elemwise_type in ["ABS", "EXP", "LOG"]: - self._op = self.__create_map[self.__elemwise_type]( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - elif self.__elemwise_type == "POW": - self._op = self.__create_map[self.__elemwise_type]( - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self._kwargs["power"], - ) - else: - self._op = self.__create_map[self.__elemwise_type]( - self.inputs[0].cnmlTensor, - self.inputs[1].cnmlTensor, - self.outputs[0].cnmlTensor, - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - if self.__elemwise_type in ["ABS", "EXP", "LOG", "POW"]: - self.__compute_map[self.__elemwise_type]( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - else: - self.__compute_map[self.__elemwise_type]( - self.op, - None, - self.inputs[0].mludata, - None, - self.inputs[1].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Reduce(OperatorBase): - __dim = None - __reduce_type = None - - __create_map = { - "AND": cnlib.cnReduceAndOp, - "OR": cnlib.cnReduceOrOp, - "SUM": cnlib.cnReduceSumOp, - "MAX": cnlib.cnReduceMaxOp, - "MEAN": cnlib.cnReduceMeanOp, - "PRODUCT": cnlib.cnReduceProductOp, - } - - __compute_map = { - "AND": cnlib.cnmlComputeReduceAndOpForward, - "OR": cnlib.cnmlComputeReduceOrOpForward, - "SUM": cnlib.cnmlComputeReduceSumOpForward_V4, - "MAX": cnlib.cnmlComputeReduceMaxOpForward_V4, - "MEAN": cnlib.cnmlComputeReduceMeanOpForward_V4, - "PRODUCT": cnlib.cnmlComputeReduceProductOpForward_V2, - } - - def __init__(self, name, inputs, outputs, dim, reduce_type): - super().__init__(name, "Reduce", inputs, outputs) - self.__dim = dim - self.__reduce_type = reduce_type - - def make_device(self): - self._op = self.__create_map[self.__reduce_type]( - self.__dim, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - self.__compute_map[self.__reduce_type]( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -# original bn in cambricon: y = (x - running_mean) * running_var -# shifted to: y = (x - running_mean) / sqrt(running_var) -class BatchNorm(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "BatchNorm", inputs, outputs) - - def _deal_param(self): - running_var = self.param_dict["running_var"].cpudata - with np.errstate(divide="ignore"): - self.param_dict["running_var"].cpudata = 1.0 / np.sqrt(running_var) - - def make_device(self): - assert "running_mean" in self.param_dict.keys(), "need mean." - assert "running_var" in self.param_dict.keys(), "need var." - self._deal_param() - self.param_dict["running_mean"].ensure_mlu() - self.param_dict["running_var"].ensure_mlu() - self._op = cnlib.cnBatchNormOp( - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["running_mean"].cnmlTensor, - self.param_dict["running_var"].cnmlTensor, - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeBatchNormOpForward_V4( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class MatMul(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "MatMul", inputs, outputs) - self._quantize_param = [] - - def make_device(self): - - self.param_dict["weight"].ensure_mlu() - - if "bias" in self.param_dict.keys(): - self.param_dict["bias"].ensure_mlu() - self._op = cnlib.cnMatMulOp( - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["weight"].cnmlTensor, - self.param_dict["bias"].cnmlTensor, - ) - else: - self._op = cnlib.cnMatMulOp( - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["weight"].cnmlTensor, - None, - ) - - input_quant = cnlib.cnQuantizedParam( - self.inputs[0]._cn_position, self.inputs[0]._cn_scale, 0 - ) - self._quantize_param.append(input_quant) - cnlib.cnmlSetOperationComputingDataType( - self._op, self.inputs[0].cnmlTensor, cnlib.CNML_DATA_INT8, input_quant - ) - if self.param_dict["weight"]._cn_data_type == cnlib.CNML_DATA_FLOAT32: - filter_quant = cnlib.cnQuantizedParam( - self.param_dict["weight"]._cn_position, - self.param_dict["weight"]._cn_scale, - 0, - ) - self._quantize_param.append(filter_quant) - cnlib.cnmlSetOperationComputingDataType( - self._op, - self.param_dict["weight"].cnmlTensor, - cnlib.CNML_DATA_INT8, - filter_quant, - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeMlpOpForward_V4( - self.op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - def destroy_param(self): - for quantize_param in self._quantize_param: - cnlib.cnDestroyQuantizedParam(quantize_param) - self._quantize_param.clear() - - -class Conv(OperatorBase): - __stride = None - __dilation = None - __pad = None - - def __init__( - self, - name, - inputs, - outputs, - stride_h, - stride_w, - dilation_h, - dilation_w, - pad_h, - pad_w, - groups=1, - ): - super().__init__(name, "Conv", inputs, outputs) - self.__stride = [stride_h, stride_w] - self.__dilation = [dilation_h, dilation_w] - self.__pad = [pad_h, pad_w] - self.groups = groups - self._quantize_param = [] - - @property - def stride(self): - return self.__stride - - @property - def pad(self): - return self.__pad - - def make_param(self): - self._param = cnlib.cnConvOpParam( - self.__stride[0], - self.__stride[1], - self.__dilation[0], - self.__dilation[1], - self.__pad[0] * 2, - self.__pad[1] * 2, - ) - - def make_device(self): - assert "W" in self.param_dict.keys(), "Need Filter." - assert "B" in self.param_dict.keys(), "Need Bias." - self.param_dict["W"].ensure_mlu() - self.param_dict["B"].ensure_mlu() - - self._ensure_param() - self._op = cnlib.cnConvOp( - self._param, - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["W"].cnmlTensor, - self.param_dict["B"].cnmlTensor, - self.groups, - ) - - input_quant = cnlib.cnQuantizedParam( - self.inputs[0]._cn_position, self.inputs[0]._cn_scale, 0 - ) - self._quantize_param.append(input_quant) - cnlib.cnmlSetOperationComputingDataType( - self._op, self.inputs[0].cnmlTensor, cnlib.CNML_DATA_INT8, input_quant - ) - - if self.param_dict["W"]._cn_data_type == cnlib.CNML_DATA_FLOAT32: - filter_quant = cnlib.cnQuantizedParam( - self.param_dict["W"]._cn_position, self.param_dict["W"]._cn_scale, 0 - ) - self._quantize_param.append(filter_quant) - cnlib.cnmlSetOperationComputingDataType( - self._op, - self.param_dict["W"].cnmlTensor, - cnlib.CNML_DATA_INT8, - filter_quant, - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - if self.groups == 1: - cnlib.cnmlComputeConvOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - else: - cnlib.cnmlComputeConvOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - def destroy_param(self): - cnlib.cnDestroyConvOpParam(self._param) - self._param = None - for quantize_param in self._quantize_param: - cnlib.cnDestroyQuantizedParam(quantize_param) - self._quantize_param.clear() - - -class Reshape(OperatorBase): - __shape = None - - def __init__(self, name, inputs, outputs): - super().__init__(name, "Reshape", inputs, outputs) - self.__shape = self.outputs[0].shape - - def make_param(self): - self._param = cnlib.cnReshapeOpParam( - self.__shape[0], - self.__shape[1], - self.__shape[2], - self.__shape[3], - self.inputs[0]._cn_tensor_layout, - ) - - def make_device(self): - self._op = cnlib.cnReshapeOp( - self._param, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeReshapeOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Pool(OperatorBase): - _map_pool_mode = { - "MAX": cnlib.CNML_POOL_MAX, - "AVG": cnlib.CNML_POOL_AVG, - } - - __windows = None - __stride = None - __pad = None - __dilation = None - __pool_mode = None - __strategy_mode = None - __real = None - - def __init__( - self, - name, - inputs, - outputs, - windows_h, - windows_w, - stride_h, - stride_w, - pad_h, - pad_w, - dilation_h, - dilation_w, - pool_mode="MAX", - strategy_mode=cnlib.CNML_POOL_KVALID, - ): - super().__init__(name, "Pool", inputs, outputs) - self.__windows = [windows_h, windows_w] - self.__stride = [stride_h, stride_w] - self.__pad = [pad_h, pad_w] - self.__dilation = [dilation_h, dilation_w] - self.__pool_mode = self._map_pool_mode[pool_mode] - self.__strategy_mode = strategy_mode - self.__real = pool_mode == "MAX" - - def make_param(self): - self._param = cnlib.cnPoolOpParam( - self.__windows[0], - self.__windows[1], - self.__stride[0], - self.__stride[1], - self.__pad[0] * 2, - self.__pad[1] * 2, - self.__dilation[0], - self.__dilation[1], - self.__pool_mode, - self.__strategy_mode, - self.__real, - ) - - def make_device(self): - self._ensure_param() - self._op = cnlib.cnPoolOp( - self._param, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputePoolOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - def destroy_param(self): - cnlib.cnDestroyPoolOpParam(self._param) - - -class Concat(OperatorBase): - _map_concat_dim = { - 0: cnlib.CNML_DIM_N, - 1: cnlib.CNML_DIM_C, - 2: cnlib.CNML_DIM_H, - 3: cnlib.CNML_DIM_W, - } - - __mode = None - - def __init__(self, name, inputs, outputs, mode=1): - super().__init__(name, "Concat", inputs, outputs) - self.__mode = self._map_concat_dim[mode] - - def make_param(self): - self._param = cnlib.cnConcatOpParam( - len(self.inputs), len(self.outputs), self.__mode - ) - - def make_device(self): - self._ensure_param() - inputs = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.inputs]) - outputs = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.outputs]) - self._op = cnlib.cnConcatOp(self._param, inputs, outputs) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - inputs = cnlib.Vectorvoid([x.mludata for x in self.inputs]) - outputs = cnlib.Vectorvoid([x.mludata for x in self.outputs]) - cnlib.cnComputeConcatOp(self._op, None, inputs, None, outputs, cnqueue, None) - - -class Slice(OperatorBase): - def __init__(self, name, inputs, outputs, slc): - super().__init__(name, "Slice", inputs, outputs) - self.nb, self.ne, self.ns = list(map(int, slc[0])) - self.cb, self.ce, self.cs = list(map(int, slc[1])) - self.hb, self.he, self.hs = list(map(int, slc[2])) - self.wb, self.we, self.ws = list(map(int, slc[3])) - - def make_param(self): - # create op param - self._param = cnlib.cnSliceOpParam( - self.nb, - self.cb, - self.hb, - self.wb, - self.ne, - self.ce, - self.he, - self.we, - self.ns, - self.cs, - self.hs, - self.ws, - ) - - def make_device(self): - # create op - self._op = cnlib.cnSliceOp( - self._param, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - # set computing data type. - # set layout - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeStridedSliceOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Broadcast(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Broadcast", inputs, outputs) - - def make_param(self): - # create op param - pass - - def make_device(self): - # create op - self._op = cnlib.cnBroadcastOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeBroadcastOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class BatchMatMul(OperatorBase): - __trans_a = None - __trans_b = None - - def __init__(self, name, inputs, outputs, trans_a=False, trans_b=False): - super().__init__(name, "BatchMatMul", inputs, outputs) - self._quantize_param = [] - self.__trans_a, self.__trans_b = trans_a, trans_b - - def make_param(self): - # create op param - pass - - def make_device(self): - # create op - self._op = cnlib.cnBatchDotOp( - self.inputs[0].cnmlTensor, - self.inputs[1].cnmlTensor, - self.outputs[0].cnmlTensor, - self.__trans_a, - self.__trans_b, - ) - cnlib.cnmlSetOperationComputingDataType( - self._op, self.inputs[0].cnmlTensor, cnlib.CNML_DATA_INT8, None - ) - cnlib.cnmlSetOperationComputingDataType( - self._op, self.inputs[1].cnmlTensor, cnlib.CNML_DATA_INT8, None - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeBatchDotOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.inputs[1].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Dimshuffle(OperatorBase): - __pattern = None - - def __init__(self, name, inputs, outputs, pattern): - super().__init__(name, "Dimshuffle", inputs, outputs) - self.__pattern = pattern - - def make_param(self): - # create op param - self._param = cnlib.cnTransposeOpParam( - self.__pattern[0], - self.__pattern[1], - self.__pattern[2], - self.__pattern[3], - self.inputs[0]._cn_tensor_layout, - ) - - def make_device(self): - # create op - self._op = cnlib.cnTransposeProOp( - self._param, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeTransposeProOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Sqrt(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Sqrt", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnSqrtOp(self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeSqrtOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Rsqrt(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Rsqrt", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnRsqrtOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeRsqrtOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Square(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Square", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnSquareOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeSquareOpForward_V2( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class BasicDiv(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "BasicDiv", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnBasicDivOp( - self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeBasicDivOpForward( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class AddPad(OperatorBase): - def __init__(self, name, inputs, outputs, pt, pb, pl, pr, *_): - super().__init__(name, "AddPad", inputs, outputs) - self._pt = pt - self._pb = pb - self._pl = pl - self._pr = pr - - def make_param(self): - self._param = cnlib.cnAddPadOpParam(self._pt, self._pb, self._pl, self._pr) - - def make_device(self): - self._op = cnlib.cnAddPadOp( - self._param, self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeAddPadOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Deconv(OperatorBase): - __stride = None - __dilation = None - __pad = None - - def __init__( - self, - name, - inputs, - outputs, - stride_h, - stride_w, - dilation_h, - dilation_w, - pad_h, - pad_w, - ): - super().__init__(name, "Deconv", inputs, outputs) - self.__stride = [stride_h, stride_w] - self.__dilation = [dilation_h, dilation_w] - self.__pad = [pad_h, pad_w] - - @property - def stride(self): - return self.__stride - - @property - def pad(self): - return self.__pad - - @property - def dilation(self): - return self.__dilation - - def make_param(self): - self._param = cnlib.cnDeconvOpParam( - self.__stride[0], - self.__stride[1], - self.__dilation[0], - self.__dilation[1], - self.__pad[0] * 2, - self.__pad[1] * 2, - ) - - def make_device(self): - assert "W" in self.param_dict.keys(), "Need Filter" - assert "B" in self.param_dict.keys(), "Need Bias" - self.param_dict["W"].ensure_mlu() - self.param_dict["B"].ensure_mlu() - - self._ensure_param() - self._op = cnlib.cnDeconvOp( - self._param, - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["W"].cnmlTensor, - self.param_dict["B"].cnmlTensor, - ) - - input_quant = cnlib.cnQuantizedParam( - self.inputs[0]._cn_position, self.inputs[0]._cn_scale, 0 - ) - cnlib.cnmlSetOperationComputingDataType( - self._op, self.inputs[0].cnmlTensor, cnlib.CNML_DATA_INT8, input_quant - ) - if self.param_dict["W"]._cn_data_type == cnlib.CNML_DATA_FLOAT32: - filter_quant = cnlib.cnQuantizedParam( - self.param_dict["W"]._cn_position, self.param_dict["W"]._cn_scale, 0 - ) - cnlib.cnmlSetOperationComputingDataType( - self._op, - self.param_dict["W"].cnmlTensor, - cnlib.CNML_DATA_INT8, - filter_quant, - ) - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeDeconvOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - def destroy_param(self): - cnlib.cnDestroyDeconvOpParam(self._param) - self._param = None - - -class ConvFirst(OperatorBase): - __stride = None - __dilation = None - __pad = None - - def __init__( - self, - name, - inputs, - outputs, - stride_h, - stride_w, - dilation_h, - dilation_w, - pad_h, - pad_w, - ): - super().__init__(name, "ConvFirst", inputs, outputs) - self.__stride = [stride_h, stride_w] - self.__dilation = [dilation_h, dilation_w] - self.__pad = [pad_h, pad_w] - - @property - def stride(self): - return self.__stride - - @property - def dilation(self): - return self.__dilation - - @property - def pad(self): - return self.__pad - - def make_param(self): - self._param = cnlib.cnConvFirstOpParam( - self.__stride[0], - self.__stride[1], - self.__dilation[0], - self.__dilation[1], - self.__pad[0], - self.__pad[1], - ) - - def make_device(self): - assert "W" in self.param_dict.keys(), "Need Filter" - assert "B" in self.param_dict.keys(), "Need Bias" - assert "mean" in self.param_dict.keys(), "Need mean" - assert "std" in self.param_dict.keys(), "Need std" - self.param_dict["W"].ensure_mlu() - self.param_dict["B"].ensure_mlu() - self.param_dict["mean"].ensure_mlu() - self.param_dict["std"].ensure_mlu() - - self._ensure_param() - self._op = cnlib.cnConvFirstOp( - self._param, - self.inputs[0].cnmlTensor, - self.outputs[0].cnmlTensor, - self.param_dict["W"].cnmlTensor, - self.param_dict["B"].cnmlTensor, - self.param_dict["mean"].cnmlTensor, - self.param_dict["std"].cnmlTensor, - ) - - cnlib.cnmlSetOperationComputingLayout( - self._op, self.inputs[0]._cn_tensor_layout - ) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeConvFirstOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - def destroy_param(self): - cnlib.cnDestroyConvFirstOpParam(self._param) - self._param = None - - -class Cast(OperatorBase): - _map_cast_type = { - "fp16-fp32": cnlib.CNML_CAST_FLOAT16_TO_FLOAT32, - "fp32-fp16": cnlib.CNML_CAST_FLOAT32_TO_FLOAT16, - } - - def __init__(self, name, inputs, outputs, cast_type): - super().__init__(name, "Cast", inputs, outputs) - self.__cast_type = Cast._map_cast_type[cast_type] - - def make_device(self): - self._op = cnlib.cnCastOp(self.inputs[0].cnmlTensor, self.outputs[0].cnmlTensor) - - def forward_trait(self, cnqueue): - cnlib.cnmlComputeCastOpForward_V4( - self._op, - None, - self.inputs[0].mludata, - None, - self.outputs[0].mludata, - cnqueue, - None, - ) - - -class Fusion(OperatorBase): - def __init__(self, name, inputs, outputs): - super().__init__(name, "Fusion", inputs, outputs) - - def make_device(self): - self._op = cnlib.cnFusionOp() - - def set_core_num(self, num): - cnlib.cnmlSetFusionOpCorenum(self._op, num) - - def fuse_op(self, base_op): - self._ensure_made() - assert not base_op.compiled, "The opr should not be compiled before fused!" - assert ( - not self._compiled - ), "Attempting fuse an opr after the fusion opr was compiled is ambitious." - base_op.make() - cnlib.cnmlFuseOp(base_op.op, self._op) - - def set_fusion_io(self): - self._ensure_made() - assert not self._compiled - inputs = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.inputs]) - outputs = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.outputs]) - cnlib.cnFusionIO(self._op, inputs, outputs) - - def add_fusion_input(self, inp): - self._ensure_made() - assert not self._compiled - cnlib.cnmlAddFusionInput(self._op, inp.cnmlTensor) - - def add_fusion_output(self, oup): - self._ensure_made() - assert not self._compiled - cnlib.cnmlAddFusionOutput(self._op, oup.cnmlTensor) - - def compile(self): - self._ensure_made() - cnlib.cnmlCompileFusionOp_V2(self._op) - self._compiled = True - - def forward_trait(self, cnqueue): - input_tensors = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.inputs]) - output_tensors = cnlib.VectorcnmlTensor([x.cnmlTensor for x in self.outputs]) - - inputs = cnlib.Vectorvoid([x.mludata for x in self.inputs]) - outputs = cnlib.Vectorvoid([x.mludata for x in self.outputs]) - - cnlib.cnComputeFusionOp( - self._op, input_tensors, inputs, output_tensors, outputs, cnqueue, None - ) - - def destroy(self): - if self._op is not None: - cnlib.cnDestroyFusionOp(self._op) - self._op = None diff --git a/mgeconvert/cambricon_converter/lib/tensor.py b/mgeconvert/cambricon_converter/lib/tensor.py deleted file mode 100644 index 2fd0e1c..0000000 --- a/mgeconvert/cambricon_converter/lib/tensor.py +++ /dev/null @@ -1,365 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import logging -from enum import Enum, unique -from functools import reduce - -import numpy as np - -from .cnlib import cambriconLib as cnlib - - -@unique -class HEAD(Enum): - UNINITIALIZED = 1 - CPU = 2 - MLU = 3 - - -@unique -class TENSOR_TYPE(Enum): - TENSOR = 0 - FILTER = 1 - CONST = 2 - VARIABLE = 3 - - -@unique -class TENSOR_LAYOUT(Enum): - NCHW = 0 - NHWC = 3 - - -@unique -class DATA_TYPE(Enum): - FLOAT16 = 1 - FLOAT32 = 2 - INT8 = 4 - UINT8 = 8 - - -class Tensor: - _map_tensor_type = { - TENSOR_TYPE.TENSOR: cnlib.CNML_TENSOR, # 0 - TENSOR_TYPE.FILTER: cnlib.CNML_FILTER, # 1 - TENSOR_TYPE.CONST: cnlib.CNML_CONST, # 2 - TENSOR_TYPE.VARIABLE: cnlib.CNML_VARIABLE, # 3 - } - - _map_cn_tensor_layout = { - TENSOR_LAYOUT.NCHW: cnlib.CNML_NCHW, # 0 - TENSOR_LAYOUT.NHWC: cnlib.CNML_NHWC, # 3 - } - - _map_sizeof = { - DATA_TYPE.INT8: 1, - DATA_TYPE.UINT8: 1, - DATA_TYPE.FLOAT16: 2, - DATA_TYPE.FLOAT32: 4, - } - - _map_np_data_type = { - DATA_TYPE.INT8: np.int8, - DATA_TYPE.UINT8: np.uint8, - # fp32 -> fp16 at cpp level - DATA_TYPE.FLOAT16: np.float32, - DATA_TYPE.FLOAT32: np.float32, - } - - _map_cn_data_type = { - DATA_TYPE.FLOAT16: cnlib.CNML_DATA_FLOAT16, # 1 - DATA_TYPE.FLOAT32: cnlib.CNML_DATA_FLOAT32, # 2 - DATA_TYPE.INT8: cnlib.CNML_DATA_INT8, # 4 - DATA_TYPE.UINT8: cnlib.CNML_DATA_UINT8, # 8 - } - - _head = None - - _cnml_tensor = None - - _tensor_type = None - _cn_tensor_type = None - - _tensor_layout = None - _cn_tensor_layout = None - - _data_type = None - _cn_data_type = None - - _data_count = 0 - _data_bytes = 0 - - _data_mlu = None - _data_cpu = None - - _scale = None - _cn_position = None - _cn_scale = None - - _name = None - - _default_dtype = DATA_TYPE.FLOAT32 - - NCHW2NHWC = False - - MUTABLE = False - - # best practice: always initilize data at the beginning - def __init__( - self, - shape, - ttype=TENSOR_TYPE.TENSOR, - tlayout=TENSOR_LAYOUT.NCHW, - dtype=None, - name=None, - data=None, - scale=None, - quantized=False, - ): - self._name = name - self.quantized = quantized - - # create cnml_tensor by tensor type - self._tensor_type = ttype - self._cn_tensor_type = self._map_tensor_type[ttype] - if self._cn_tensor_type is not cnlib.CNML_VARIABLE: - self._cnml_tensor = cnlib.cnTensor_V2(self._cn_tensor_type) - if Tensor.MUTABLE and ttype == TENSOR_TYPE.TENSOR: - cnlib.cnmlSetTensorDimMutable( - self._cnml_tensor, [True, False, False, False] - ) - else: - self._cnml_tensor = cnlib.cnTensor_V3() - cnlib.cnmlSetTensorType(self._cnml_tensor, cnlib.CNML_VARIABLE) - - # set cnml_tensor data type - if dtype is None: - dtype = Tensor._default_dtype - if dtype not in self._map_cn_data_type.keys(): - raise NotImplementedError( - "host data type {} was not supported now.".format(self.dtype) - ) - self._data_type = dtype - self._cn_data_type = self._map_cn_data_type[dtype] - cnlib.cnmlSetTensorDataType(self._cnml_tensor, self._cn_data_type) - - if Tensor.NCHW2NHWC: - tlayout = TENSOR_LAYOUT.NHWC - # set layout - if tlayout not in self._map_cn_tensor_layout.keys(): - raise NotImplementedError( - "host tensor layout {} was not supported now.".format(tlayout) - ) - self._tensor_layout = tlayout - self._cn_tensor_layout = self._map_cn_tensor_layout[tlayout] - - # set cnml_tensor shape - shape = list(shape) - if len(shape) < 4: - shape += [1] * (4 - len(shape)) - self._data_shape = shape - if Tensor.NCHW2NHWC: - cnlib.cnmlSetTensorShape( - self._cnml_tensor, self._get_internal_shape(self._data_shape) - ) - else: - cnlib.cnmlSetTensorShape(self._cnml_tensor, self._data_shape) - - # set cnml_tensor position and scale - if scale is not None: - self._set_scale(scale) - - if self._cn_data_type == cnlib.CNML_DATA_INT8: - assert scale is not None, "quantized tensor must set scale" - cnlib.cnmlSetQuantizedPosition(self._cnml_tensor, self._cn_position) - cnlib.cnmlSetQuantizedScale(self._cnml_tensor, self._cn_scale) - - # _data_mlu: cpu->mlu for inputs, mlu->cpu for outputs - self._data_count = reduce(lambda x, y: x * y, self._data_shape, 1) - self._data_bytes = self._data_count * self._map_sizeof[self._data_type] - if self._cn_tensor_type is cnlib.CNML_TENSOR: - self._data_mlu = cnlib.cnMalloc(self._data_bytes) - - if self._cn_tensor_type in (cnlib.CNML_FILTER, cnlib.CNML_CONST): - assert data is not None, "data must be specified for filter/const tensor" - - if data is None: - self._head = HEAD.UNINITIALIZED - else: - self._set_data_cpu(data) - - def _get_quantized_info(self, scale): - assert isinstance(scale, (float, int)) - if scale < 2 ** -32 or scale > 2 ** 32: - logging.warning( - "scale is not in the valid range:(%s, %s), which will result in " - "incorrect calculation of the converted model.", - 2 ** -32, - 2 ** 32, - ) - summ = 1.0 - cn_position = 0 - if scale >= 1.0: - while summ < scale: - cn_position = cn_position + 1 - summ = summ * 2 - elif scale > 0.0: - while summ >= scale: - cn_position = cn_position - 1 - summ = summ * 0.5 - cn_position = cn_position + 1 - summ = summ * 2 - else: - raise ValueError("scale {} should > 0".format(scale)) - - if cn_position > 32: - cn_position = 32 - if cn_position < -32: - cn_position = -32 - - cn_scale = float(2 ** cn_position) / scale - - return cn_position, cn_scale - - # NCHW -> NHWC - def _get_internal_shape(self, shape): - internal_shape = list(shape) - internal_shape[1], internal_shape[2], internal_shape[3] = ( - internal_shape[2], - internal_shape[3], - internal_shape[1], - ) - return internal_shape - - @property - def mludata(self): - return self._data_mlu - - @property - def cpudata(self): - if self._data_cpu is None: - return None - if Tensor.NCHW2NHWC: - ret = self._data_cpu.astype(self._map_np_data_type[self._data_type]) - ret = ret.reshape(self._get_internal_shape(self._data_shape)) - ret = ret.transpose((0, 3, 1, 2)) - else: - ret = self._data_cpu.reshape(self._data_shape) - return ret - - @cpudata.setter - def cpudata(self, data): - if data is None: - self._data_cpu = None - return - - if self._cn_tensor_type in (cnlib.CNML_CONST, cnlib.CNML_FILTER): - return - - self._set_data_cpu(data) - - def _set_data_cpu(self, data): - if isinstance(data, (int, float)): - data = [data] * self._data_count - data = ( - np.asarray(data) - .astype(self._map_np_data_type[self._data_type]) - .reshape(self._data_shape) - ) - if Tensor.NCHW2NHWC: - data = data.transpose((0, 2, 3, 1)) - data = data.flatten() - - self._data_cpu = data - self._head = HEAD.CPU - - @property - def scale(self): - return self._scale - - @scale.setter - def scale(self, scale): - if scale is None: - self._scale = None - self._cn_position = None - self._cn_scale = None - self._set_scale(scale) - - def _set_scale(self, scale): - self._scale = scale - self._cn_position, self._cn_scale = self._get_quantized_info(self._scale) - - @property - def head(self): - return self._head - - @property - def shape(self): - return self._data_shape - - @property - def name(self): - return self._name - - @property - def ttype(self): - return self._tensor_type - - @property - def tlayout(self): - return self._tensor_layout - - @property - def dtype(self): - return self._data_type - - @property - def cnmlTensor(self): - return self._cnml_tensor - - def h2d(self): - if self._data_cpu is not None: - if self._cn_tensor_type in (cnlib.CNML_FILTER, cnlib.CNML_CONST): - if self._cn_data_type is cnlib.CNML_DATA_INT8: - cnlib.cnH2dConstInt8(self._cnml_tensor, self._data_cpu, False) - elif self._cn_data_type is cnlib.CNML_DATA_FLOAT16: - cnlib.cnH2dConstFloat16(self._cnml_tensor, self._data_cpu, False) - else: - cnlib.cnH2dConst(self._cnml_tensor, self._data_cpu, False) - else: - # cpu->mlu: for inputs - if self._cn_data_type is cnlib.CNML_DATA_UINT8: - cnlib.cnH2dUint8(self._data_mlu, self._data_cpu) - elif self._cn_data_type is cnlib.CNML_DATA_FLOAT16: - cnlib.cnH2dFloat16(self._data_mlu, self._data_cpu) - else: - cnlib.cnH2d(self._data_mlu, self._data_cpu) - self._head = HEAD.MLU - - def d2h(self): - if self._data_mlu is not None: - self._head = HEAD.MLU - if self._data_cpu is None: - self._data_cpu = np.array([0.0] * self._data_count).astype( - self._map_np_data_type[self._data_type] - ) - # mlu->cpu: for outputs - if self._cn_data_type is cnlib.CNML_DATA_FLOAT16: - cnlib.cnD2hFloat16(self._data_cpu, self._data_mlu) - else: - cnlib.cnD2h(self._data_cpu, self._data_mlu) - - def ensure_mlu(self): - if self._head != HEAD.MLU: - self.h2d() - - def destroy(self): - if self._data_mlu is not None: - cnlib.cnrtFree(self._data_mlu) - cnlib.cnDestroyTensor(self._cnml_tensor) diff --git a/mgeconvert/cambricon_converter/swig/cambricon.i b/mgeconvert/cambricon_converter/swig/cambricon.i deleted file mode 100644 index 573b692..0000000 --- a/mgeconvert/cambricon_converter/swig/cambricon.i +++ /dev/null @@ -1,529 +0,0 @@ -/* - * $File: cambricon.i - * - * $Copyright: Copyright (c) 2014-2020 Megvii Inc. All rights reserved. - */ - -%module cambriconLib - -%{ -#define SWIG_FILE_WITH_INIT -#include "cnrt.h" -#include "cnml.h" -%} - -// ignore functions those are not implemented -%ignore cnmlSetBaseOpJobType; -%ignore cnmlGetBitWidth; -%ignore cnmlSetFusionOpCoreVersionChangable; -%ignore cnmlGetOffset; -%ignore cnmlGetScale; -%ignore cnmlComputeNdCycleNEqualOpForward; -%ignore cnmlGetPosition; -%ignore cnmlSetFusionOperationComputingLayout; -%ignore cnmlCreateNdCycleNEqualOp; -%ignore cnmlGetMovingPosition; -%ignore cnmlComputeNdCycleGreaterOpForward; -%ignore cnmlCreateCastOpForward; -%ignore cnmlCreateBroadcastMultOpForward; -%ignore cnmlGetInterval; -%ignore cnmlComputeSquaredDiffOpForward_V3; -%ignore cnmlSetBasicDivHighPrecision; -%ignore cnmlCreateBroadcastOpForward; -%ignore cnmlCreateArgminOp; - -// use numpy swig for python-c type convertions, mainly (python list/numpy array <==> c array) -%include "numpy.i" -%init %{ - import_array(); -%} - -%apply (int DIM1, int *IN_ARRAY1) { (int dim_nums, int dim_values[]) }; -%apply (int DIM1, float *IN_ARRAY1) { (int data_count, float *data_cpu) }; -%apply (int DIM1, int8_t *IN_ARRAY1) { (int data_count, int8_t *data_cpu) }; -%apply (int DIM1, uint8_t *IN_ARRAY1) { (int data_count, uint8_t* data_cpu) }; -%apply (int DIM1, int16_t *IN_ARRAY1) { (int data_count, int16_t* data_cpu) }; -%apply (int DIM1, float *INPLACE_ARRAY1) { (int data_out_count, float *data_out_cpu) }; -%apply (int DIM1, int16_t *INPLACE_ARRAY1) { (int data_out_count, int16_t *data_out_cpu) }; -%apply (int DIM1, int *INPLACE_ARRAY1) { (int dim, int *shape) }; - -// cnmlStatus_t cnmlSetTensorDimMutable(cnmlTensor_t tensor, bool *dim_mutable, int dim_num); -%typemap(in) (bool *dim_mutable, int dim_num) { - if (!PyList_Check($input)) { - PyErr_SetString(PyExc_ValueError, "Expecting a list"); - return NULL; - } - $2 = PyList_Size($input); - $1 = (bool *) malloc(($2) * sizeof(bool)); - int i; - for (i = 0; i < $2; ++i) { - PyObject *s = PyList_GetItem($input, i); - if (!PyBool_Check(s)) { - free($1); - PyErr_SetString(PyExc_ValueError, "List items must be bools"); - return NULL; - } - // $1[i] = PyBool_AsBool(s); - // $1[i] = PyNumber_AsSsize_t(s); - $1[i] = (s == Py_True); - } -} -%typemap(freearg) (bool *dim_mutable, int dim_num) { - if ($1) free($1); -} - -// void cnGetShape(cnmlTensor_t tensor, int tensor_shape[4]) -// will return a list with 4 elements in Python -%typemap(in, numinputs=0) int tensor_shape[4] (int temp[4]) { - $1 = temp; -} - -%typemap(argout) int tensor_shape[4] { - int i; - $result = PyList_New(4); - for (i = 0; i < 4; i++) { - PyObject *o = PyInt_FromLong((int) $1[i]); - PyList_SetItem($result, i, o); - } -} - -%include "std_vector.i" - -%include "cnrt.h" -%include "cnml.h" - -// handle convertion between array of class and python list -namespace std { - %template(VectorcnmlTensor) vector; - %template(Vectorvoid) vector; -}; - -%inline{ - cnmlTensor_t cnTensor_V2(cnmlTensorType_t tensor_type) { - cnmlTensor_t ptensor; - cnmlCreateTensor_V2(&ptensor, tensor_type); - return ptensor; - } - - cnmlTensor_t cnTensor_V3() { - cnmlTensor_t ptensor; - cnmlCreateTensor_V3(&ptensor); - return ptensor; - } - - void cnDestroyTensor(cnmlTensor_t ptensor) { - cnmlDestroyTensor(&ptensor); - } - - cnmlModel_t cnModel(char *name) { - cnmlModel_t p_model; - cnmlCreateModel(&p_model, name); - return p_model; - } - - cnrtQueue_t cnQueue() { - cnrtQueue_t pqueue; - cnrtCreateQueue(&pqueue); - return pqueue; - } - - void *cnMalloc(size_t sz) { - void *pvar; - cnrtRet_t ret = cnrtMalloc(&pvar, sz); - return pvar; - } - - void cnH2d(void *data_mlu, int data_count, float *data_cpu) { - void *void_data_cpu = reinterpret_cast(data_cpu); - cnrtMemcpy(data_mlu, void_data_cpu, data_count * sizeof(float), - CNRT_MEM_TRANS_DIR_HOST2DEV); - } - - void cnH2dUint8(void *data_mlu, int data_count, uint8_t* data_cpu) { - void* void_data_cpu = reinterpret_cast(data_cpu); - cnrtMemcpy(data_mlu, void_data_cpu, data_count * sizeof(uint8_t), - CNRT_MEM_TRANS_DIR_HOST2DEV); - } - - void cnH2dFloat16(void* data_mlu, int data_count, float* data_cpu) { - // input h2d: cpu fp32 -> cpu int16 -> dev fp16 - // declare cpu int16 - int16_t* int16_data_cpu = (int16_t*)malloc(data_count * sizeof(int16_t)); - // cpu fp32 -> cpu int16 - cnrtCastDataType(data_cpu, CNRT_FLOAT32, int16_data_cpu, CNRT_FLOAT16, data_count, NULL); - // cpu int16 -> dev fp16 - cnrtMemcpy(data_mlu, int16_data_cpu, data_count * sizeof(int16_t), - CNRT_MEM_TRANS_DIR_HOST2DEV); - } - - void cnD2h(int data_out_count, float *data_out_cpu, void *data_mlu) { - void *tmp_data_cpu = reinterpret_cast(data_out_cpu); - cnrtMemcpy(tmp_data_cpu, data_mlu, data_out_count * sizeof(float), - CNRT_MEM_TRANS_DIR_DEV2HOST); - data_out_cpu = reinterpret_cast(tmp_data_cpu); - } - - void cnD2hFloat16(int data_out_count, float* data_out_cpu, void* data_mlu) { - // output d2h: dev fp16 -> cpu int16 -> cpu fp32 - // declare cpu int16 - int16_t* int16_data_cpu = (int16_t*)malloc(data_out_count * sizeof(int16_t)); - // dev fp16 -> cpu int16 - cnrtMemcpy(int16_data_cpu, data_mlu, data_out_count * sizeof(int16_t), - CNRT_MEM_TRANS_DIR_DEV2HOST); - // cpu int16 -> cpu fp32 - cnrtCastDataType(int16_data_cpu, CNRT_FLOAT16, data_out_cpu, CNRT_FLOAT32, data_out_count, NULL); - - } - - void cnH2dConstInt8(cnmlTensor_t tensor, int data_count, int8_t *data_cpu, bool free_aftercompile) { - void *void_data_cpu = reinterpret_cast(data_cpu); - cnmlBindConstData_V2(tensor, void_data_cpu, free_aftercompile); - } - - void cnH2dConstFloat16(cnmlTensor_t tensor, int data_count, float *data_cpu, bool free_aftercompile) { - // bind const fp16: float32 -> float16 -> bind - int16_t* int16_data_cpu = (int16_t*)malloc(data_count * sizeof(int16_t)); - cnrtCastDataType(data_cpu, CNRT_FLOAT32, int16_data_cpu, CNRT_FLOAT16, data_count, NULL); - cnmlBindConstData_V2(tensor, int16_data_cpu, free_aftercompile); - } - - void cnH2dConst(cnmlTensor_t tensor, int data_count, float *data_cpu, bool free_aftercompile) { - void *void_data_cpu = reinterpret_cast(data_cpu); - cnmlBindConstData_V2(tensor, void_data_cpu, free_aftercompile); - } - - cnmlBaseOp_t cnCastOp(cnmlTensor_t inp, cnmlTensor_t oup, cnmlCastType_t type) { - cnmlBaseOp_t op; - cnmlCreateCastOp(&op, type, inp, oup); - return op; - } - - cnmlFusionOp_t cnFusionOp() { - cnmlFusionOp_t fuse_op; - cnmlCreateFusionOp(&fuse_op); - return fuse_op; - } - - void cnGetShape(cnmlTensor_t tensor, int tensor_shape[4]) { - cnmlGetTensorShape(tensor, tensor_shape); - } - - cnmlBaseOp_t cnDevMemcpyOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t cpy_op; - cnmlCreateDeviceMemcpyOp(&cpy_op, inp, oup); - return cpy_op; - } - -#define REG_ELEMWISE_1(type) \ - cnmlBaseOp_t cn##type##Op(cnmlTensor_t inp, cnmlTensor_t oup) \ - { \ - cnmlBaseOp_t elemwise_op; \ - cnmlCreate##type##Op(&elemwise_op, inp, oup); \ - return elemwise_op; \ - } -REG_ELEMWISE_1(Abs) -REG_ELEMWISE_1(Exp) -REG_ELEMWISE_1(Log) -#undef REG_ELEMWISE - - cnmlBaseOp_t cnBasicDivOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t basic_div_op; - cnmlCreateBasicDivOp(&basic_div_op, inp, oup); - return basic_div_op; - } - - cnmlBaseOp_t cnSqrtOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t sqrt_op; - cnmlCreateSqrtOp(&sqrt_op, inp, oup); - return sqrt_op; - } - - cnmlBaseOp_t cnRsqrtOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t rsqrt_op; - cnmlCreateRsqrtOp(&rsqrt_op, inp, oup); - return rsqrt_op; - } - - cnmlBaseOp_t cnSquareOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t square_op; - cnmlCreateSquareOp(&square_op, inp, oup); - return square_op; - } - - cnmlBaseOp_t cnPowOp(cnmlTensor_t inp, cnmlTensor_t oup, float power_c) { - cnmlBaseOp_t pow_op; - cnmlCreatePowerOp(&pow_op, inp, oup, power_c); - return pow_op; - } - - cnmlBaseOp_t cnArgmaxOp(cnmlDimension_t dim_axis, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t argmax_op; - cnmlCreateArgmaxOp(&argmax_op, dim_axis, inp, oup); - return argmax_op; - } - - cnmlBaseOp_t cnFloorOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t floor_op; - cnmlCreateFloorOp(&floor_op, inp, oup); - return floor_op; - } - - cnmlBaseOp_t cnActiveOp(cnmlActiveFunction_t active_func, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t active_op; - cnmlCreateActiveOp(&active_op, active_func, inp, oup); - return active_op; - } - -#define REG_ELEMWISE_2(type) \ - cnmlBaseOp_t cn##type##Op(cnmlTensor_t inp1, cnmlTensor_t inp2, cnmlTensor_t oup) \ - { \ - cnmlBaseOp_t elemwise_op; \ - cnmlCreateBroadcast##type##Op(&elemwise_op, inp1, inp2, oup); \ - return elemwise_op; \ - } -REG_ELEMWISE_2(Add) -REG_ELEMWISE_2(Sub) -REG_ELEMWISE_2(Mult) -#undef REG_ELEMWISE - - cnmlBaseOp_t cnDivOp(cnmlTensor_t inp1, cnmlTensor_t inp2, cnmlTensor_t oup) { - cnmlBaseOp_t div_op; - cnmlCreateRealDivOp(&div_op, inp1, inp2, oup); - return div_op; - } - -#define REG_CYCLE(type) \ - cnmlBaseOp_t cnCycle##type##Op(cnmlTensor_t inp1, cnmlTensor_t inp2, cnmlTensor_t oup) { \ - cnmlBaseOp_t op; \ - cnmlCreateCycle##type##Op(&op, inp1, inp2, oup); \ - return op; \ - } -REG_CYCLE(Add) -REG_CYCLE(Mult) -#undef REG_CYCLE - -#define REG_REDUCE(type) \ - cnmlBaseOp_t cnReduce##type##Op(cnmlDimension_t dim, cnmlTensor_t inp, cnmlTensor_t oup) \ - { \ - cnmlBaseOp_t reduce_op; \ - cnmlCreateReduce##type##Op(&reduce_op, dim, inp, oup); \ - return reduce_op; \ - } -REG_REDUCE(Max) -REG_REDUCE(Mean) -REG_REDUCE(Sum) -REG_REDUCE(Product) -#undef REG_REDUCE - - cnmlBaseOp_t cnReduceAndOp(cnmlReduce_andDim_t dim, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t reduce_op; - cnmlCreateReduceAndOp(&reduce_op, dim, inp, oup); - return reduce_op; - } - - cnmlBaseOp_t cnReduceOrOp(cnmlReduce_orDim_t dim, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t reduce_op; - cnmlCreateReduceOrOp(&reduce_op, dim, inp, oup); - return reduce_op; - } - - cnmlBaseOp_t cnBatchNormOp(cnmlTensor_t inp, cnmlTensor_t oup, cnmlTensor_t mean, cnmlTensor_t var) { - cnmlBaseOp_t bn_op; - cnmlCreateBatchNormOpForward(&bn_op, inp, oup, mean, var); - return bn_op; - } - - cnmlBaseOp_t cnMatMulOp(cnmlTensor_t inp, cnmlTensor_t oup, cnmlTensor_t filter, cnmlTensor_t bias) { - cnmlBaseOp_t matmul_op; - cnmlCreateMlpOp(&matmul_op, inp, oup, filter, bias); - return matmul_op; - } - - cnmlBaseOp_t cnBatchDotOp(cnmlTensor_t lef, cnmlTensor_t rht, cnmlTensor_t oup, bool trans_a, bool trans_b) { - cnmlBaseOp_t batchdot_op; - cnmlCreateBatchDotOp(&batchdot_op, lef, rht, oup, trans_a, trans_b); - return batchdot_op; - } - - cnmlConvOpParam_t cnConvOpParam(int stride_h, int stride_w, - int dilation_h, int dilation_w, - int pad_h, int pad_w) { - cnmlConvOpParam_t param; - cnmlCreateConvOpParam(¶m, stride_h, stride_w, dilation_h, dilation_w, pad_h, pad_w); - return param; - } - - void cnDestroyConvOpParam(cnmlConvOpParam_t param) { - cnmlDestroyConvOpParam(¶m); - } - - cnmlDeconvOpParam_t cnDeconvOpParam(int stride_h, int stride_w, - int dilation_h, int dilation_w, - int pad_h, int pad_w) { - cnmlDeconvOpParam_t param; - cnmlCreateDeconvOpParam_V3(¶m, stride_h, stride_w, 0, 0, 0, 0, pad_w, pad_h, dilation_h, dilation_w); - return param; - } - - void cnDestoryDeconvOpParam(cnmlDeconvOpParam_t param) { - cnmlDestroyDeconvOpParam(¶m); - } - - cnmlQuantizedParam_t cnQuantizedParam(int pos, float scale = 1., float offset = 0.) { - cnmlQuantizedParam_t input_quant_param; - cnmlCreateQuantizedParam(&input_quant_param, pos, scale, offset); - return input_quant_param; - } - - void cnDestroyQuantizedParam(cnmlQuantizedParam_t qparam) { - cnmlDestroyQuantizedParam(&qparam); - } - - cnmlBaseOp_t cnConvOp(cnmlConvOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup, - cnmlTensor_t filter, cnmlTensor_t bias, int groups) { - cnmlBaseOp_t conv_op; - if (groups == 1) cnmlCreateConvOpForward(&conv_op, param, inp, oup, filter, bias); - else cnmlCreateConvGroupOpForward(&conv_op, param, inp, oup, filter, bias, groups); - return conv_op; - } - - cnmlBaseOp_t cnDeconvOp(cnmlDeconvOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup, - cnmlTensor_t filter, cnmlTensor_t bias) { - cnmlBaseOp_t deconv_op; - cnmlCreateDeconvOpForward(&deconv_op, param, inp, oup, filter, bias); - return deconv_op; - } - - cnmlConvFirstOpParam_t cnConvFirstOpParam(int sh, int sw, int dh, int dw, int ph, int pw) { - cnmlConvFirstOpParam_t param; - cnmlCreateConvFirstOpParam_V2(¶m, sh, sw, dh, dw, pw, pw, ph, ph); - return param; - } - - cnmlBaseOp_t cnConvFirstOp(cnmlConvFirstOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup, - cnmlTensor_t filter, cnmlTensor_t bias, - cnmlTensor_t mean, cnmlTensor_t std) { - cnmlBaseOp_t op; - cnmlCreateConvFirstOpForward(&op, param, inp, mean, oup, filter, bias, std); - return op; - } - - void cnDestroyConvFirstOpParam(cnmlConvFirstOpParam_t param) { - cnmlDestroyConvFirstOpParam(¶m); - } - - cnmlReshapeOpParam_t cnReshapeOpParam(int no, int co, int ho, int wo, - cnmlDataOrder_t data_order = CNML_NHWC) { - cnmlReshapeOpParam_t param; - cnmlCreateReshapeOpParam(¶m, no, co, ho, wo, data_order); - return param; - } - - cnmlBaseOp_t cnReshapeOp(cnmlReshapeOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t reshape_op; - cnmlCreateReshapeOp(&reshape_op, param, inp, oup); - return reshape_op; - } - - cnmlAddPadOpParam_t cnAddPadOpParam(int pt, int pb, int pl, int pr, float pad_value) { - cnmlAddPadOpParam_t param; - cnmlCreateAddPadOpParam_V2(¶m, pt, pb, pl, pr, pad_value); - return param; - } - - cnmlBaseOp_t cnAddPadOp(cnmlAddPadOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t op; - cnmlCreateAddPadOp(&op, param, inp, oup); - return op; - } - - cnmlStridedSliceOpParam_t cnSliceOpParam(int nb, int cb, int hb, int wb, int ne, int ce, int he, int we, int ns, int cs, int hs, int ws) { - cnmlStridedSliceOpParam_t param; - cnmlCreateStridedSliceOpParam(¶m, nb, cb, hb, wb, ne, ce, he, we, ns, cs, hs, ws); - return param; - } - - cnmlBaseOp_t cnSliceOp(cnmlStridedSliceOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t slice_op; - cnmlCreateStridedSliceOp(&slice_op, param, inp, oup); - return slice_op; - } - - cnmlTransposeOpParam_t cnTransposeOpParam(int d0, int d1, int d2, int d3, - cnmlDataOrder_t data_order = CNML_NCHW) { - cnmlTransposeOpParam_t param; - cnmlCreateTransposeOpParam(¶m, data_order, d0, d1, d2, d3); - return param; - } - - cnmlBaseOp_t cnTransposeProOp(cnmlTransposeOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t transpose_op; - cnmlCreateTransposeProOp(&transpose_op, inp, oup, param); - return transpose_op; - } - - cnmlBaseOp_t cnBroadcastOp(cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t broadcast_op; - cnmlCreateBroadcastOp(&broadcast_op, inp, oup); - return broadcast_op; - } - - cnmlPoolOpParam_t cnPoolOpParam(int window_h, int window_w, int stride_h, - int stride_w, int pad_h, int pad_w, int dilation_h, int dilation_w, - cnmlPoolMode_t pool_mode = CNML_POOL_MAX, - cnmlPoolStrategyMode_t strategy_mode = CNML_POOL_KVALID, - bool real = false) { - cnmlPoolOpParam_t param; - cnmlCreatePoolOpParam(¶m, window_h, window_w, stride_h, stride_w, - pad_h, pad_w, dilation_h, dilation_w, pool_mode, strategy_mode, real); - return param; - } - - void cnDestroyPoolOpParam(cnmlPoolOpParam_t param) { - cnmlDestroyPoolOpParam(¶m); - } - - cnmlBaseOp_t cnPoolOp(cnmlPoolOpParam_t param, cnmlTensor_t inp, cnmlTensor_t oup) { - cnmlBaseOp_t pool_op; - cnmlCreatePoolOp(&pool_op, param, inp, oup); - return pool_op; - } - - cnmlConcatOpParam_t cnConcatOpParam(int input_num, int output_num, cnmlDimension_t concat_mode) { - cnmlConcatOpParam_t param; - cnmlCreateConcatOpParam(¶m, input_num, output_num, concat_mode); - return param; - } - - cnmlBaseOp_t cnConcatOp(cnmlConcatOpParam_t param, - std::vector input_tensors, - std::vector output_tensors) { - cnmlBaseOp_t concat_op; - cnmlCreateConcatOp(&concat_op, param, input_tensors.data(), input_tensors.size(), output_tensors.data(), output_tensors.size()); - return concat_op; - } - - void cnComputeConcatOp(cnmlBaseOp_t concatop, cnmlTensor_t *inp_tensors, std::vector inputs, - cnmlTensor_t *oup_tensors, std::vector outputs, cnrtQueue_t cnq, void *extra) { - cnmlComputeConcatOpForward_V4(concatop, inp_tensors, inputs.data(), inputs.size(), oup_tensors, outputs.data(), outputs.size(), cnq, extra); - } - - void cnFusionIO(cnmlFusionOp_t op, std::vector input_tensors, std::vector output_tensors) { - cnmlSetFusionIO(op, input_tensors.data(), input_tensors.size(), output_tensors.data(), output_tensors.size()); - } - - // void cnComputeFusionOp(cnmlFusionOp_t op, cnmlTensor_t input_tensors[], std::vector inputs, - // cnmlTensor_t output_tensors[], std::vector outputs, cnrtQueue_t cnq, void *extra) { - void cnComputeFusionOp(cnmlFusionOp_t op, std::vector input_tensors, std::vector inputs, - std::vector output_tensors, std::vector outputs, cnrtQueue_t cnq, void* extra) { - cnmlComputeFusionOpForward_V4(op, input_tensors.data(), inputs.data(), inputs.size(), output_tensors.data(), outputs.data(), outputs.size(), cnq, extra); - } - - void cnDestroyBaseOp(cnmlBaseOp_t op) { - cnmlDestroyBaseOp(&op); - } - - void cnDestroyFusionOp(cnmlFusionOp_t op) { - cnmlDestroyFusionOp(&op); - } -} diff --git a/mgeconvert/utils/__init__.py b/mgeconvert/converter_ir/__init__.py similarity index 93% rename from mgeconvert/utils/__init__.py rename to mgeconvert/converter_ir/__init__.py index 1207b5d..93515f3 100644 --- a/mgeconvert/utils/__init__.py +++ b/mgeconvert/converter_ir/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. diff --git a/mgeconvert/converter_ir/ir_graph.py b/mgeconvert/converter_ir/ir_graph.py new file mode 100644 index 0000000..a5b9225 --- /dev/null +++ b/mgeconvert/converter_ir/ir_graph.py @@ -0,0 +1,157 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from typing import List # pylint: disable=unused-import +from typing import Sequence + +from megengine.logger import get_logger + +from .ir_op import OpBase +from .ir_tensor import IRTensor + +logger = get_logger(__name__) + + +class IRGraph: + def __init__(self) -> None: + self.graph_inputs = [] # type: List[IRTensor] + self.graph_outputs = [] # type: List[IRTensor] + + self._tensor_ids = [] # type: List[int] + self._opr_ids = [] # type: List[int] + self.all_tensors = [] # type: List[IRTensor] + self.all_oprs = [] # type: List[OpBase] + + def add_op(self, op, index=None): + assert len(self._opr_ids) == len(self.all_oprs) + if index: + self._opr_ids.insert(index, id(op)) + self.all_oprs.insert(index, op) + else: + if not isinstance(op, Sequence): + ops = (op,) + for opr in ops: + self.all_oprs.append(opr) + self._opr_ids.append(id(opr)) + + def delete_ops(self, index): + del self.all_oprs[index] + del self._opr_ids[index] + + def get_tensor(self, var_id, ir_tensor: IRTensor, origin_tensor: IRTensor = None): + """ + var_id: origin var id + origin_tensor: relate the var to a already existed ir tensor + """ + if var_id: + if var_id not in self._tensor_ids: + self._tensor_ids.append(var_id) + if origin_tensor: + self.all_tensors.append(origin_tensor) + else: + self.all_tensors.append(ir_tensor) + return self.all_tensors[self._tensor_ids.index(var_id)] + else: # int, float + return ir_tensor + + def add_tensor(self, var_id, ir_tensor): + if var_id not in self._tensor_ids: + self._tensor_ids.append(var_id) + self.all_tensors.append(ir_tensor) + + def find_inp_oprs(self, op): + if len(op.inp_tensors) == 0: + return None + inp_oprs = [] + for inp in op.inp_tensors: + if inp.owner_opr is None: + continue + inp_oprs.append(inp.owner_opr) + return inp_oprs + + def find_out_oprs(self, op): + out_oprs = [] + for oup in op.out_tensors: + if oup.user_opr is not None: + out_oprs.extend(oup.user_opr) + return out_oprs + + def add_net_inputs(self, inp_tensor: IRTensor): + self.graph_inputs.append(inp_tensor) + inp_tensor.owner_opr = self + + def add_net_outputs(self, out_tensor: IRTensor): + self.graph_outputs.append(out_tensor) + + def insert_op_after(self, new_op: OpBase, last_op: OpBase): + """ + only consider cases of one output + """ + out_oprs = self.find_out_oprs(last_op) + if len(out_oprs) == 0: + # last op of the graph + for t in new_op.inp_tensors: + t.user_opr.append(new_op) + if t in self.graph_outputs: + self.graph_outputs.remove(t) + for t in new_op.out_tensors: + t.owner_opr = new_op + self.graph_outputs.append(t) + else: + assert len(out_oprs) == 1, "Do not support more than one output oprs yet." + next_op = out_oprs[0] + for t in new_op.inp_tensors: + if t.owner_opr == last_op: + assert next_op in t.user_opr + idx = t.user_opr.index(next_op) + t.user_opr[idx] = new_op + for t in new_op.out_tensors: + t.owner_opr = new_op + assert t in next_op.inp_tensors + t.user_opr.append(next_op) + op_idx = self.all_oprs.index(last_op) + 1 + self.add_op(new_op, op_idx) + + def insert_op_before(self, new_op: OpBase, next_op: OpBase): + inp_oprs = self.find_inp_oprs(next_op) + if len(inp_oprs) == 0: + # first op of the graph + for t in new_op.inp_tensors: + assert next_op in t.user_opr + idx = t.user_opr.index(next_op) + t.user_opr[idx] = new_op + if t not in self.graph_inputs: + self.graph_inputs.append(t) + logger.warning("WARNING: Add graph inputs(%s)", t.name) + for t in new_op.out_tensors: + t.owner_opr = new_op + else: + assert len(inp_oprs) == 1, "Do not support more than one input oprs yet." + last_op = inp_oprs[0] + for t in new_op.inp_tensors: + if t.owner_opr == last_op: + assert next_op in t.user_opr + idx = t.user_opr.index(next_op) + t.user_opr[idx] = new_op + for t in new_op.out_tensors: + t.owner_opr = new_op + assert t in next_op.inp_tensors + t.user_opr.append(next_op) + op_idx = self.all_oprs.index(next_op) + self.add_op(new_op, op_idx) + + def __repr__(self): + res = "" + for op in self.all_oprs: + res += op.name + "\n" + res += "\tInput:\n" + for i in op.inp_tensors: + res += "\t\t{}\n".format(i.name) + res += "\tOutput:\n" + for i in op.out_tensors: + res += "\t\t{}\n".format(i.name) + return res diff --git a/mgeconvert/converter_ir/ir_op.py b/mgeconvert/converter_ir/ir_op.py new file mode 100644 index 0000000..098a04d --- /dev/null +++ b/mgeconvert/converter_ir/ir_op.py @@ -0,0 +1,401 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from typing import List # pylint: disable=unused-import + +from .ir_tensor import IRTensor # pylint: disable=unused-import + + +class OpBase: + skip = False + name = "" + + def __init__(self) -> None: + self.inp_tensors = [] # type: List[IRTensor] + self.out_tensors = [] # type: List[IRTensor] + self.activation = "IDENTITY" + + def add_inp_tensors(self, ir_tensor): + self.inp_tensors.append(ir_tensor) + + def add_out_tensors(self, ir_tensor): + self.out_tensors.append(ir_tensor) + + +#################### conv related ############################ +class _ConvOpr(OpBase): + def __init__(self, stride, padding, dilation, groups): + super().__init__() + self.stride = stride + self.padding = padding + self.dilation = dilation + self.groups = groups + + +class Conv2dOpr(_ConvOpr): + name = "Conv2d" + + +class Deconv2dOpr(_ConvOpr): + name = "Deconv2d" + + +class ConvolutionBackwardFilterOpr(OpBase): + name = "ConvolutionBackwardFilter" + + def __init__( + self, stride, padding, dilation, group, kernel_shape, src_shape, grad_out_shape + ): + super().__init__() + self.stride = stride + self.padding = padding + self.dilation = dilation + self.group = group + self.kernel_shape = kernel_shape + self.src_shape = src_shape + self.grad_out_shape = grad_out_shape + + +class _PoolOpr(OpBase): + def __init__(self, kernel_size, stride, padding): + super().__init__() + self.kernel_size = kernel_size + self.stride = stride + self.padding = padding + + +class MaxPool2dOpr(_PoolOpr): + name = "MaxPool2d" + + def __init__(self, kernel_size, stride, padding): + super().__init__(kernel_size, stride, padding) + self.mode = "MAX" + + +class AvgPool2dOpr(_PoolOpr): + name = "AvgPool2d" + + def __init__( + self, kernel_size, stride, padding, mode="AVERAGE_COUNT_EXCLUDE_PADDING" + ): + super().__init__(kernel_size, stride, padding) + self.mode = mode + + +class PadOpr(OpBase): + name = "Pad" + + +class BatchNormalizationOpr(OpBase): + name = "BatchNormalization" + + def __init__(self, output_idx=0): + super().__init__() + self.output_idx = output_idx + + +class AdaptiveAvgPool2dOpr(OpBase): + name = "AdaptiveAvgPool2d" + + def __init__(self, out_shape): + super().__init__() + self.out_shape = out_shape + + +#################### math related ############################ + + +class MatMulOpr(OpBase): + name = "MatMul" + + def __init__( + self, + transpose_a=False, + transpose_b=False, + compute_mode="default", + format="default", + ): + super().__init__() + self.transpose_a = transpose_a + self.transpose_b = transpose_b + self.compute_mode = compute_mode + self.format = format + + +class LinearOpr(MatMulOpr): + name = "Linear" + + def __init__(self, has_bias=False): + super().__init__(transpose_b=True) + self.has_bias = has_bias + + +class ReduceOpr(OpBase): + name = "Reduce" + + def __init__(self, axis, mode, keep_dims): + super().__init__() + self.axis = axis + self.mode = mode + self.keep_dims = keep_dims + + +class SoftmaxOpr(OpBase): + name = "Softmax" + + def __init__(self, axis=None, beta=1): + super().__init__() + self.axis = axis + self.beta = beta + + +#################### tensor related ############################ +class FlattenOpr(OpBase): + name = "Flatten" + + def __init__(self, start_axis=0, end_axis=-1): + super().__init__() + self.start_axis = start_axis + self.end_axis = end_axis + + +class DropoutOpr(OpBase): + name = "Dropout" + + def __init__(self, drop_prob=0, training=False): + super().__init__() + self.drop_prob = drop_prob + self.training = training + + +class ConstantOpr(OpBase): + name = "Constant" + + +class MultipleDeviceTensorHolderOpr(OpBase): + name = "MultipleDeviceTensorHolder" + + +class SharedDeviceTensorOpr(OpBase): + name = "SharedDeviceTensorOpr" + + +class VolatileSharedDeviceTensorOpr(OpBase): + name = "VolatileSharedDeviceTensor" + + +class GetVarShapeOpr(OpBase): + name = "GetVarShape" + + +class IndexingOneHotOpr(OpBase): + name = "IndexingOneHotOpr" + + +class LinspaceOpr(OpBase): + name = "Linspace" + + +class WarpPerspectiveForwardOpr(OpBase): + name = "WarpPerspectiveForward" + + +class IdentityOpr(OpBase): + name = "Identity" + + def __init__(self): + super().__init__() + self.mode = "Identity" + + +class ConcatOpr(OpBase): + name = "Concat" + + def __init__(self, axis): + super().__init__() + self.axis = axis + + +class ReshapeOpr(OpBase): + name = "Reshape" + + def __init__(self, out_shape): + super().__init__() + self.out_shape = out_shape + + +class TransposeOpr(OpBase): + name = "Transpose" + + def __init__(self, pattern: list): + super().__init__() + self.pattern = pattern + + +class SqueezeOpr(OpBase): + name = "Squeeze" + + def __init__(self, squeeze_dims): + super().__init__() + self.squeeze_dims = squeeze_dims + + +class GetSubTensorOpr(OpBase): + name = "GetSubTensor" + + def __init__(self, axis, begin_params, end_params, step_params, squeeze_axis=None): + super().__init__() + self.axis = axis + self.begin_params = begin_params + self.end_params = end_params + self.step_params = step_params + self.squeeze_axis = squeeze_axis + + +class ResizeOpr(OpBase): + name = "Resize" + + def __init__( + self, out_size, scale_factor=None, mode="bilinear", align_corners=None + ): + super().__init__() + self.out_size = out_size + self.scale_factor = scale_factor + self.mode = mode + self.align_corners = align_corners + + +class AxisAddRemoveOpr(OpBase): + name = "AxisAddRemove" + + def __init__(self, out_shape, desc): + super().__init__() + self.out_shape = out_shape + self.desc = desc + + +class BroadcastOpr(OpBase): + name = "Broadcast" + + +############################ elemwise ######################## + + +class ElemwiseOpr(OpBase): + pass + + +class AddOpr(ElemwiseOpr): + name = "Add" + + +class SubOpr(ElemwiseOpr): + name = "Sub" + + +class MulOpr(ElemwiseOpr): + name = "Mul" + + +class TrueDivOpr(ElemwiseOpr): + name = "TrueDiv" + + +class PowOpr(ElemwiseOpr): + name = "Pow" + + +class ExpOpr(ElemwiseOpr): + name = "Exp" + + +class FloorOpr(ElemwiseOpr): + name = "Floor" + + +class FloorDivOpr(ElemwiseOpr): + name = "FloorDiv" + + +class CeilOpr(ElemwiseOpr): + name = "Ceil" + + +class MaxOpr(ElemwiseOpr): + name = "Max" + + +class MinOpr(ElemwiseOpr): + name = "Min" + + +class AbsOpr(ElemwiseOpr): + name = "Abs" + + +class LogOpr(ElemwiseOpr): + name = "Log" + + +class FuseMulAdd3Opr(OpBase): + name = "FuseMulAdd3" + + +############################# activation ########################### + + +class Relu6Opr(OpBase): + name = "Relu6" + + +class ReluOpr(OpBase): + name = "Relu" + + +class SigmoidOpr(OpBase): + name = "Sigmoid" + + +class HardSigmoidOpr(OpBase): + name = "HardSigmoid" + + +class SiLUOpr(OpBase): + name = "SiLU" + + +class TanHOpr(OpBase): + name = "TanH" + + +class LeakyReluOpr(OpBase): + name = "LeakyRelu" + + def __init__(self, negative_slope=0.01): + super().__init__() + self.negative_slope = negative_slope + + +class TypeCvtOpr(OpBase): + name = "TypeCvt" + + def __init__(self, out_dtype): + super().__init__() + self.out_dtype = out_dtype + + +class HardSwishOpr(OpBase): + name = "HardSwish" + + +class RepeatOpr(OpBase): + name = "Repeat" + + def __init__(self, repeats, axis): + super().__init__() + self.repeats = repeats + self.axis = 0 if axis is None else axis diff --git a/mgeconvert/converter_ir/ir_quantizer.py b/mgeconvert/converter_ir/ir_quantizer.py new file mode 100644 index 0000000..d87e62b --- /dev/null +++ b/mgeconvert/converter_ir/ir_quantizer.py @@ -0,0 +1,86 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint:disable=import-outside-toplevel, no-name-in-module +import json + +import megengine +import numpy as np + +from .ir_tensor import IRTensor + + +class IRQuantizer: + def __init__(self, require_quantize=False, param_fake_quant=False): + super().__init__() + self.require_quantize = require_quantize + self.param_fake_quant = param_fake_quant + + def quantize(self, tensor: IRTensor): + assert self.require_quantize, "This net do not require true quantize." + value = tensor.np_data + if isinstance(value, megengine.tensor): + value = value.numpy() + if tensor.scale: + value = value / tensor.scale + value = np.round(value) + if tensor.zero_point: + value += tensor.zero_point + dt = ( + np.dtype(tensor.q_dtype) + if isinstance(tensor.q_dtype, str) + else tensor.q_dtype + ) + if np.issubdtype(dt, np.integer): + v_min = np.iinfo(dt).min + v_max = np.iinfo(dt).max + value = np.clip(value, v_min, v_max) + value = value.astype(tensor.q_dtype) + return value + + def save_quantize_params(self, irgraph, path="quant_params.json"): + quant_params = {} + all_tensors = set() + for opr in irgraph.all_oprs: + for t in opr.inp_tensors + opr.out_tensors: + all_tensors.add(t) + for t in all_tensors: + dt = np.dtype(t.q_dtype) + v_max, v_min = None, None + is_weight = bool(t.np_data is not None) + if np.issubdtype(dt, np.integer): + v_min = np.iinfo(dt).min + v_max = np.iinfo(dt).max + if self.param_fake_quant and is_weight: + if t.scale is not None: + inp = megengine.tensor(t.np_data) + scale = megengine.tensor(t.scale) + zp = float(t.zero_point) if t.zero_point else 0.0 + zero_point = megengine.tensor(zp) + from megengine.core._imperative_rt.core2 import ( # pylint:disable=import-error + apply, + ) + from megengine.core.ops.builtin import FakeQuant + + t.np_data = apply( + FakeQuant(qmin=v_min, qmax=v_max), inp, scale, zero_point + )[0].numpy() + else: + param = { + "dtype": str(dt), + "qmin": str(v_min), + "qmax": str(v_max), + "scale": str(t.scale), + "zero_point": str(t.zero_point), + "is_weight": is_weight, + } + quant_params[t.name] = param + + params = json.dumps(quant_params, indent=4) + with open(path, "w") as f: + f.write(params) diff --git a/mgeconvert/converter_ir/ir_tensor.py b/mgeconvert/converter_ir/ir_tensor.py new file mode 100644 index 0000000..44081f5 --- /dev/null +++ b/mgeconvert/converter_ir/ir_tensor.py @@ -0,0 +1,141 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from typing import List, Sequence, Union + + +class DataFormat: + @staticmethod + def permute_shape(shape, order): + return [shape[i] for i in order] + + +class NHWCFormat: + def __init__(self) -> None: + self.permute_to_NCHW = [0, 3, 1, 2] + + def shape_to_NCHW(self, data): + assert isinstance(data.axis_order, NHWCFormat) + assert data.ndim == 4 + return DataFormat.permute_shape(data.shape, self.permute_to_NCHW) + + def data_to_NCHW(self, data): + assert data.ndim == 4 + return data.transpose(self.permute_to_NCHW) + + +class NCHWFormat: + def __init__(self): + self.permute_to_NHWC = [0, 2, 3, 1] + + def shape_to_NHWC(self, shape): + assert len(shape) == 4 + return DataFormat.permute_shape(shape, self.permute_to_NHWC) + + def data_to_NHWC(self, data): + assert data.ndim == 4 + return data.transpose(self.permute_to_NHWC) + + +class IOHWFormat: + def __init__(self): + self.permute_to_OHWI = [1, 2, 3, 0] + + def shape_to_OHWI(self, shape): + assert len(shape) == 4 + return DataFormat.permute_shape(shape, self.permute_to_OHWI) + + def data_to_OHWI(self, data): + assert data.ndim == 4 + return data.transpose(self.permute_to_OHWI) + + +class OIHWFormat(NCHWFormat): + def shape_to_OHWI(self, shape): + return super().shape_to_NHWC(shape) + + def shape_to_IHWO(self, shape): + return [shape[1], shape[2], shape[3], shape[0]] + + def data_to_OHWI(self, data): + return super().data_to_NHWC(data) + + def data_to_IHWO(self, data): + return data.transpose(1, 2, 3, 0) + + +class OHWIFormat(NHWCFormat): + def shape_to_OIHW(self, data): + return super().shape_to_NCHW(data) + + def data_to_OIHW(self, data): + return super().data_to_NCHW(data) + + +class AxisOrder: + NCHW = NCHWFormat() + NHWC = NHWCFormat() + + OIHW = OIHWFormat() + OHWI = OHWIFormat() + + IOHW = IOHWFormat() # deconv weight + + +class IRTensor: + def __init__( + self, + name, + shape, + dtype, + scale=None, + zero_point=None, + q_type=None, + np_data=None, + owner_opr=None, + axis=AxisOrder.NCHW, + ): + self.name = name + self.axis_order = axis + self.owner_opr = owner_opr + self.user_opr = [] + self.shape = shape + self.dtype = dtype + + self.np_data = np_data + + self.scale = scale + self.zero_point = zero_point + self.q_dtype = q_type + + @property + def ndim(self): + return len(self.shape) + + def add_user_opr(self, op): + self.user_opr.append(op) + + def set_dtype(self, target_type): + self.np_data = self.np_data.astype(target_type) + self.dtype = target_type + + def set_qparams( + self, scale: Union[float, List[float]], zero_point=None, q_dtype=None + ): + if not isinstance(scale, Sequence): # per tensor + self.scale = float(scale) + else: # per channel + self.scale = [float(s) for s in scale] + + if zero_point is not None: + if not isinstance(zero_point, Sequence): + self.zero_point = int(zero_point) + else: + self.zero_point = [int(zp) for zp in zero_point] + + if self.q_dtype is not None: + self.q_dtype = q_dtype diff --git a/mgeconvert/converter_ir/ir_transform.py b/mgeconvert/converter_ir/ir_transform.py new file mode 100644 index 0000000..150cfb2 --- /dev/null +++ b/mgeconvert/converter_ir/ir_transform.py @@ -0,0 +1,890 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from collections import OrderedDict +from enum import Enum +from functools import cmp_to_key +from typing import Set # pylint: disable=unused-import +from typing import Callable, Dict, Sequence + +import numpy as np + +from ..converter_ir.ir_graph import IRGraph +from .ir_op import ( + AddOpr, + Conv2dOpr, + Deconv2dOpr, + DropoutOpr, + ExpOpr, + FuseMulAdd3Opr, + GetSubTensorOpr, + HardSigmoidOpr, + HardSwishOpr, + LeakyReluOpr, + MulOpr, + OpBase, + PadOpr, + ReduceOpr, + ReluOpr, + ReshapeOpr, + ResizeOpr, + SoftmaxOpr, + SqueezeOpr, + SubOpr, + TanHOpr, + TransposeOpr, + TrueDivOpr, + _PoolOpr, +) +from .ir_tensor import AxisOrder, IRTensor + + +class TransformerRule(Enum): + # general rules + NOPE = 1 + + # for TFLite + REDUCE_AXIS_AS_INPUT = 100 + REMOVE_RESHAPE_INPUT = 101 + # FUSE_FOR_RELU6 pass should happen before FUSE_ACTIVATION + FUSE_FOR_RELU6 = 102 ## + FUSE_ACTIVATION = 103 + CONV_ADD_ZERO_BIAS = 104 + DECONV_ADD_ZERO_BIAS = 105 + # DEPTHWISE_CONV_RESHAPE_WEIGHT requirs RESHAPE_BIAS_TO_1DIM + DEPTHWISE_CONV_RESHAPE_WEIGHT = 106 + FUSE_SOFTMAX = 107 + # RESHAPE_BIAS_TO_1DIM should happen before DECONV_SHAPE_AS_INPUT + RESHAPE_BIAS_TO_1DIM = 108 + DECONV_SHAPE_AS_INPUT = 109 + FUSE_ASTYPE = 110 ## + PADDING_FOR_CONV_AND_POOLING = 111 + TRANSPOSE_PATTERN_AS_INPUT = 112 + # FUSE_FOR_LEAKY_RELU should happen before EXPAND_MUL_ADD3 + FUSE_FOR_LEAKY_RELU = 113 + EXPAND_MUL_ADD3 = 114 + EXPAND_ADD_SIGMOID = 115 ## + FUSE_FOR_CONV_BIAS = 116 + FUSE_FOR_DECONV_BIAS = 117 + FUSE_FOR_FULLY_CONNECTED = 118 ## + # for TFLite Converter + SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE = 119 + RESIZE_PARAMS_AS_INPUT = 120 + + # remove reshape + REMOVE_RESHAPE_REALTED_OP = 121 + REMOVE_DROPOUT = 122 + REMOVE_UNRELATED_IROP = 130 + ADD_FAKE_HSIGMOID_OUT = 131 + + +def cmp_rules(a, b): + if a.value < b.value: + return -1 + if a.value > b.value: + return 1 + return 0 + + +class IRTransform: + def __init__(self, transformer_options): + if not isinstance(transformer_options, Sequence): + transformer_options = [ + transformer_options, + ] + + # bias of depthwise_conv must be 1 dim + if TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT in transformer_options: + if TransformerRule.RESHAPE_BIAS_TO_1DIM not in transformer_options: + transformer_options.append(TransformerRule.RESHAPE_BIAS_TO_1DIM) + + self.trans_options = sorted(transformer_options, key=cmp_to_key(cmp_rules)) + + def transform(self, ir_graph): + for option in self.trans_options: + TRANSFORMMAP[option](ir_graph) + return ir_graph + + +TRANSFORMMAP: Dict[Enum, Callable] = {} + + +def _register_tranformation_rule(transformer_option): + def callback(impl): + TRANSFORMMAP[transformer_option] = impl + + return callback + + +def cal_pad_mode(tm_opr): + out_shape = tm_opr.out_tensors[0].shape + inp_shape = tm_opr.inp_tensors[0].shape + if out_shape[2:] == inp_shape[2:]: + return "SAME" + else: + return "VALID" + + +@_register_tranformation_rule(TransformerRule.REMOVE_RESHAPE_INPUT) +def _remove_reshape_input(net): + for op in net.all_oprs: + if not isinstance(op, ReshapeOpr): + continue + + if len(op.inp_tensors) == 2: + del op.inp_tensors[1] + # TODO： delete _tensor_ids, all_tensors + + +@_register_tranformation_rule(TransformerRule.TRANSPOSE_PATTERN_AS_INPUT) +def _transpose_pattern_as_input(net): + for op in net.all_oprs: + if not isinstance(op, TransposeOpr): + continue + + perm_tensor = IRTensor( + name=op.inp_tensors[0].name + "_perm", + shape=np.array(op.pattern).shape, + dtype=np.int32, + np_data=np.array(op.pattern, dtype=np.int32), + owner_opr=op, + q_type=np.int32, + axis=None, + ) + op.add_inp_tensors(perm_tensor) + + +@_register_tranformation_rule(TransformerRule.REDUCE_AXIS_AS_INPUT) +def _reduce_axis_as_input(net): + for op in net.all_oprs: + if not isinstance(op, ReduceOpr): + continue + + axis_tensor = IRTensor( + name=op.inp_tensors[0].name + "_axis", + shape=[1], + dtype=np.int32, + np_data=np.array(op.axis, dtype=np.int32), + owner_opr=op, + q_type=np.int32, + axis=None, + ) + op.add_inp_tensors(axis_tensor) + + +@_register_tranformation_rule(TransformerRule.PADDING_FOR_CONV_AND_POOLING) +def _make_padding(net: IRGraph): + def have_padding(opr): + if isinstance(opr, Conv2dOpr): + if cal_pad_mode(opr) == "SAME": + return False + if hasattr(opr, "padding") and (opr.padding[0] > 0 or opr.padding[1] > 0): + return True + return False + + insert_intended = OrderedDict() # type: OrderedDict + for op in net.all_oprs: + if not isinstance(op, (Conv2dOpr, _PoolOpr)): + continue + + if have_padding(op): + assert op.inp_tensors[0].ndim == 4, "ERROR: unsupported padding mode" + np_data = np.array( + [ + 0, + 0, + op.padding[0], + op.padding[0], + op.padding[1], + op.padding[1], + 0, + 0, + ], + dtype=np.int32, + ) + + new_tensor_id = max(net._tensor_ids) + 1 + pad_in_tensor = IRTensor( + name=op.inp_tensors[0].name + "_paddings", + shape=[4, 2], + dtype=np.int32, + owner_opr=None, + np_data=np_data, + q_type=np.int32, + axis=None, + ) + net.add_tensor(new_tensor_id, pad_in_tensor) + + shape = list(op.inp_tensors[0].shape) + new_tensor_id = max(net._tensor_ids) + 1 + pad_out_tensor = IRTensor( + name=op.inp_tensors[0].name + "_pad_out", + shape=[ + shape[0], + shape[1], + shape[2] + op.padding[0] * 2, + shape[3] + op.padding[1] * 2, + ], + dtype=op.inp_tensors[0].dtype, + ) + if ( + hasattr(op.inp_tensors[0], "scale") + and op.inp_tensors[0].scale is not None + ): + pad_out_tensor.scale = op.inp_tensors[0].scale + pad_out_tensor.q_dtype = op.inp_tensors[0].q_dtype + if hasattr(op.inp_tensors[0], "zero_point"): + pad_out_tensor.zero_point = op.inp_tensors[0].zero_point + net.add_tensor(new_tensor_id, pad_out_tensor) + + pad_opr = PadOpr() + pad_opr.inp_tensors = [op.inp_tensors[0], pad_in_tensor] + index = op.inp_tensors[0].user_opr.index(op) + op.inp_tensors[0].user_opr[index] = pad_opr + pad_opr.out_tensors = [pad_out_tensor] + pad_out_tensor.owner_opr = pad_opr + op.inp_tensors = [pad_out_tensor] + op.inp_tensors[1:] + pad_out_tensor.user_opr.append(op) + + index = net._opr_ids.index(id(op)) + insert_intended[index] = (id(pad_opr), pad_opr) + + for index, generated_pair in list(insert_intended.items())[::-1]: + net._opr_ids.insert(index, generated_pair[0]) + net.all_oprs.insert(index, generated_pair[1]) + + +@_register_tranformation_rule(TransformerRule.DECONV_SHAPE_AS_INPUT) +def _deconv_shape_as_input(net: IRGraph): + for op in net.all_oprs: + if not isinstance(op, Deconv2dOpr): + continue + + result_shape = op.out_tensors[0].shape + np_data = np.array( + [result_shape[0], result_shape[2], result_shape[3], result_shape[1],], + dtype=np.int32, + ) + new_tensor_id = max(net._tensor_ids) + 1 + shape_symvar = IRTensor( + name=op.inp_tensors[0].name + "_deconv_out_shape", + shape=[4], + dtype=np.int32, + owner_opr=op, + np_data=np_data, + q_type=np.int32, + axis=None, + ) + shape_tensor = net.get_tensor(new_tensor_id, shape_symvar) + if len(op.inp_tensors) == 2: + op.inp_tensors = [ + shape_tensor, + op.inp_tensors[1], + op.inp_tensors[0], + ] + else: + op.inp_tensors = [ + shape_tensor, + op.inp_tensors[1], + op.inp_tensors[0], + op.inp_tensors[2], + ] + + +@_register_tranformation_rule(TransformerRule.RESIZE_PARAMS_AS_INPUT) +def _resize_params_as_input(net): + for op in net.all_oprs: + if not isinstance(op, ResizeOpr): + continue + + if len(op.inp_tensors) == 2: + continue + + out_size_tensor = IRTensor( + name=op.inp_tensors[0].name + "_out_size", + shape=(2,), + dtype=np.int32, + np_data=np.array(op.out_size, dtype=np.int32), + q_type=np.int32, + axis=None, + ) + op.add_inp_tensors(out_size_tensor) + + +@_register_tranformation_rule(TransformerRule.CONV_ADD_ZERO_BIAS) +def _add_bias_for_conv(net: IRGraph): + for op in net.all_oprs: + if not isinstance(op, Conv2dOpr): + continue + if len(op.inp_tensors) == 3: + continue + + weight_shape = op.inp_tensors[1].shape + bias_shape = ( + weight_shape[0] + if len(weight_shape) == 4 + else weight_shape[0] * weight_shape[1] + ) + bias_shape = (1, bias_shape, 1, 1) + bias = np.zeros(bias_shape, dtype=np.float32) + bias_tensor = IRTensor( + name=op.inp_tensors[0].name + "_bias", + shape=bias_shape, + dtype=np.float32, + np_data=bias, + axis=AxisOrder.NCHW, + ) + if op.inp_tensors[0].scale and op.inp_tensors[1].scale: + bias_tensor.set_qparams( + op.inp_tensors[0].scale * op.inp_tensors[1].scale, 0 + ) + bias_tensor.q_dtype = "int32" + op.inp_tensors.append(bias_tensor) + + +@_register_tranformation_rule(TransformerRule.DECONV_ADD_ZERO_BIAS) +def _add_bias_for_deconv(net: IRGraph): + for op in net.all_oprs: + if not isinstance(op, Deconv2dOpr): + continue + if len(op.inp_tensors) == 3: + continue + + weight_shape = op.inp_tensors[1].shape + bias_shape = ( + weight_shape[1] + if len(weight_shape) == 4 + else weight_shape[0] * weight_shape[2] + ) + bias_shape = (1, bias_shape, 1, 1) + bias = np.zeros(bias_shape, dtype=np.float32) + bias_tensor = IRTensor( + name=op.inp_tensors[0].name + "_bias", + shape=bias_shape, + dtype=np.float32, + np_data=bias, + axis=AxisOrder.NCHW, + ) + if op.inp_tensors[0].scale and op.inp_tensors[1].scale: + bias_tensor.set_qparams( + op.inp_tensors[0].scale * op.inp_tensors[1].scale, 0 + ) + bias_tensor.q_dtype = "int32" + op.inp_tensors.append(bias_tensor) + + +@_register_tranformation_rule(TransformerRule.RESHAPE_BIAS_TO_1DIM) +def _reshape_bias_to_1dim(net: IRGraph): + for op in net.all_oprs: + if not isinstance(op, (Deconv2dOpr, Conv2dOpr)): + continue + if len(op.inp_tensors) == 2: + continue + + bias = op.inp_tensors[2] + if bias.ndim == 4: + bias.shape = (bias.shape[1],) + bias.np_data = bias.np_data.reshape(-1) + + +@_register_tranformation_rule(TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT) +def _depthwise_conv_reshape_weight(net: IRGraph): + # general group conv is not supported for TFLite + for op in net.all_oprs: + if not isinstance(op, Conv2dOpr): + continue + if op.groups == 1: + continue + + weight = op.inp_tensors[1] # G, oc/G, ic/G, kh, kw + ic, cm = weight.shape[1] * op.groups, weight.shape[2] + h, w = weight.shape[3:5] + weight.shape = (ic, cm, h, w) # oc, ic/G, kh, kw + weight.np_data = weight.np_data.reshape(ic, cm, h, w) + + +@_register_tranformation_rule(TransformerRule.FUSE_ACTIVATION) +def _fuse_activation(net): + delete_intended = [] + + for op_id, op in zip(net._opr_ids, net.all_oprs): + if isinstance(op, (ReluOpr, TanHOpr)): + prev_ops = net.find_inp_oprs(op) + if len(prev_ops) == 0: + continue + prev_op = prev_ops[0] + if not isinstance(prev_op, OpBase): + continue + if prev_op.activation != "IDENTITY" or prev_op.name == "Deconv2d": + continue + + activation = op.name.upper() + prev_op.activation = activation + prev_op.out_tensors = op.out_tensors + + delete_intended.append(net._opr_ids.index(op_id)) + + for delete_idx in delete_intended[::-1]: + net.delete_ops(delete_idx) + + +@_register_tranformation_rule(TransformerRule.SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE) +def _make_slice_as_inputs(net: IRGraph): + for op in net.all_oprs: + if not isinstance(op, GetSubTensorOpr): + continue + + ndim = op.inp_tensors[0].ndim + + def make_input(axis, param, init_value): + # make inputs: begin, end and step. + ret = [init_value] * ndim # pylint:disable=cell-var-from-loop + for k, v in zip(axis, param): + ret[k] = v + ret = IRTensor( + name=op.name + "_fake_input", # pylint:disable=cell-var-from-loop + shape=[len(ret)], + dtype=np.int32, + np_data=np.array(ret, dtype=np.int32), + owner_opr=op, # pylint:disable=cell-var-from-loop + q_type=np.int32, + ) + return ret + + begins_tensor = make_input(op.axis, op.begin_params, 0) + ends_tensor = make_input(op.axis, op.end_params, np.iinfo(np.int32).max) + steps_tensor = make_input(op.axis, op.step_params, 1) + + op.inp_tensors = [op.inp_tensors[0], begins_tensor, ends_tensor, steps_tensor] + + # TFLite slice do not support squeeze axis, so insert a squeeze opr here. + # infer actual output shape of tflite slice + desired_out_shape = op.out_tensors[0].shape + actual_out_shape = [1] * ndim + idx = 0 + for i in range(ndim): + if i in op.squeeze_axis: + continue + actual_out_shape[i] = desired_out_shape[idx] + idx += 1 + slice_out_tensor = IRTensor( + name=op.name + "fake_output", + shape=actual_out_shape, + dtype=op.out_tensors[0].dtype, + q_type=op.out_tensors[0].q_dtype, + owner_opr=op, + ) + old_out = op.out_tensors + op.out_tensors = [slice_out_tensor] + + squeeze = SqueezeOpr(op.squeeze_axis) + squeeze.inp_tensors = [slice_out_tensor] + squeeze.out_tensors = old_out + + idx = net._opr_ids.index(id(op)) + 1 + net.add_op(squeeze, idx) + + +# caffe transormer rules +class PatternNode: + def __init__(self, type, is_output=False, const_value=None): + self.op = None + self.type = type + self.inp_oprs = [] + self.inp_const = [] + self.inp_tensors = [] + self.is_output = is_output + self.const_value = const_value + + def check_const_value(self, op): + inp_tensors = [v.np_data for v in op.inp_tensors] + for const in self.const_value: + idx = const[0] + if idx == -1: + find = False + for index, v in enumerate(inp_tensors): + if np.array_equal(const[1], v): + find = True + del inp_tensors[index] + break + if not find: + return False + elif not np.array_equal(const[1], inp_tensors[idx]): + return False + return True + + +get_type = lambda op: type(op).__name__ + + +def match(node, opr): + node_queue = [node] + opr_queue = [opr] + matched_opr = set() + matched_node = set() + while len(node_queue) != 0: + cur_node = node_queue.pop(0) + cur_opr = opr_queue.pop(0) + if cur_node.type != get_type(cur_opr) and cur_node.type != "*" or cur_opr.skip: + return False + if cur_node.op == None: + cur_node.op = cur_opr + if cur_node.const_value != None: + if not cur_node.check_const_value(cur_opr): + return False + elif cur_node.op != cur_opr: + return False + + matched_opr.add(cur_opr) + matched_node.add(cur_node) + for i, var in enumerate(cur_opr.inp_tensors): + if var.np_data is not None: + cur_node.inp_const.append([i, var.np_data]) + else: + cur_node.inp_tensors.append([i, var]) + if len(cur_node.inp_oprs) == 0: + continue + if len(cur_node.inp_oprs) != len(cur_opr.inp_oprs): + return False + + for i, j in zip(cur_node.inp_oprs, cur_opr.inp_oprs): + node_queue.append(i) + opr_queue.append(j) + + for n in matched_node: + if n.is_output: + continue + for op in n.op.out_oprs: + if op not in matched_opr: + return False + + return True + + +def get_softmax_axis(ndim: int) -> int: + if ndim in (0, 1, 3): + return 0 + return 1 + + +@_register_tranformation_rule(TransformerRule.FUSE_SOFTMAX) +def _fuse_softmax(net: IRGraph): + matches = OrderedDict() # type: OrderedDict + + for op in net.all_oprs: + if not isinstance(op, TrueDivOpr): + continue + try: + prev_op = net.find_inp_oprs(op)[1] + cur_index = net._opr_ids.index(id(op)) + if ( + not isinstance(prev_op, ReduceOpr) + or prev_op.mode != "SUM" + or prev_op.axis != get_softmax_axis(prev_op.inp_tensors[0].ndim) + or net._opr_ids.index(id(prev_op)) != cur_index - 1 + ): + continue + prev_op = net.find_inp_oprs(op)[0] + if ( + not isinstance(prev_op, ExpOpr) + or net._opr_ids.index(id(prev_op)) != cur_index - 2 + ): + continue + prev_op = net.find_inp_oprs(prev_op)[0] + if ( + not isinstance(prev_op, SubOpr) + or net._opr_ids.index(id(prev_op)) != cur_index - 3 + ): + continue + + prev_op = net.find_inp_oprs(prev_op)[1] + if ( + not isinstance(prev_op, ReduceOpr) + or prev_op.mode != "MAX" + or prev_op.axis != get_softmax_axis(prev_op.inp_tensors[0].ndim) + or net._opr_ids.index(id(prev_op)) != cur_index - 4 + ): + continue + except IndexError: # doesn't match + continue + + softmax_opr = SoftmaxOpr(axis=get_softmax_axis(prev_op.inp_tensors[0].ndim)) + softmax_opr.beta = 1 + softmax_opr.inp_tensors = prev_op.inp_tensors[:1] + for i in softmax_opr.inp_tensors: + i.user_opr.append(softmax_opr) + softmax_opr.out_tensors = op.out_tensors + softmax_out_oprs = net.find_out_oprs(op) + matches[id(prev_op)] = (id(prev_op), softmax_opr, softmax_out_oprs) + + for original_id, generated_pair in list(matches.items())[::-1]: + index = net._opr_ids.index(original_id) + for out_op in generated_pair[2]: + generated_pair[1].out_tensors[0].user_opr.append(out_op) + + del net._opr_ids[index : index + 5] + del net.all_oprs[index : index + 5] + + net._opr_ids.insert(index, generated_pair[0]) + net.all_oprs.insert(index, generated_pair[1]) + + +@_register_tranformation_rule(TransformerRule.FUSE_FOR_LEAKY_RELU) +def _fuse_leaky_relu(net: IRGraph): + """ + Elemwise(ADD) + Elemwise(MUL) + Elemwise(MAX) + Elemwise(MIN) -> LeakyRelu + """ + for opr in net.all_oprs: + if ( + opr.name == "Add" + and len(net.find_inp_oprs(opr)) == 2 + and net.find_inp_oprs(opr)[0].name == "Max" + and net.find_inp_oprs(opr)[1].name == "Mul" + ): + max_op = net.find_inp_oprs(opr)[0] + mul_op = net.find_inp_oprs(opr)[1] + if not mul_op.inp_tensors[1].shape == (1,): + continue + if not max_op.inp_tensors[1].shape == (1,): + continue + if ( + len(net.find_inp_oprs(mul_op)) != 1 + or net.find_inp_oprs(mul_op)[0].name != "Min" + or net.find_inp_oprs(mul_op)[0].inp_tensors[1].shape != (1,) + ): + continue + min_op = net.find_inp_oprs(mul_op)[0] + if not min_op.inp_tensors[1].shape == (1,): + continue + if max_op.inp_tensors[0] != min_op.inp_tensors[0]: + continue + leaky_relu = LeakyReluOpr( + negative_slope=float(mul_op.inp_tensors[1].np_data) + ) + leaky_relu.inp_tensors = [max_op.inp_tensors[0]] + max_op.inp_tensors[0].user_opr.remove(max_op) + max_op.inp_tensors[0].user_opr.remove(min_op) + max_op.inp_tensors[0].user_opr.append(leaky_relu) + leaky_relu.out_tensors = opr.out_tensors + opr.out_tensors[0].owner_opr = leaky_relu + + index = net.all_oprs.index(max_op) + del net.all_oprs[index : index + 4] + del net._opr_ids[index : index + 4] + net.add_op(leaky_relu, index) + + +@_register_tranformation_rule(TransformerRule.FUSE_FOR_CONV_BIAS) +def _fuse_for_conv_bias(net: IRGraph): + """ + ConvolutionForward + Elemwise(ADD) -> ConvForwardBias + """ + for opr in net.all_oprs: + if ( + opr.name == "Conv2d" + and len(net.find_out_oprs(opr)) == 1 + and net.find_out_oprs(opr)[0].name == "Add" + ): + bias_op = net.find_out_oprs(opr)[0] + if not ( + ( + bias_op.inp_tensors[1].np_data is not None + and len(bias_op.inp_tensors[1].np_data.reshape(-1)) + == opr.inp_tensors[1].shape[0] + ) + or ( + ( + bias_op.inp_tensors[0].np_data is not None + and len(bias_op.inp_tensors[0].np_data.reshape(-1)) + == opr.inp_tensors[1].shape[0] + ) + ) + ): + continue + bias_idx = 0 if bias_op.inp_tensors[0].np_data is not None else 1 + if len(opr.inp_tensors) == 2: + opr.inp_tensors.append(bias_op.inp_tensors[bias_idx]) + else: + opr.inp_tensors[2].np_data += bias_op.inp_tensors[ + bias_idx + ].np_data.reshape(-1) + if bias_op in opr.out_tensors[0].user_opr: + opr.out_tensors[0].user_opr.remove(bias_op) + bias_out_op = net.find_out_oprs(bias_op) + if len(bias_out_op) > 0: + for op in bias_out_op: + op.inp_tensors[0] = opr.out_tensors[0] + opr.out_tensors[0].user_opr.append(op) + else: + # last op of the graph + assert bias_op.out_tensors[0] in net.graph_outputs + index = net.graph_outputs.index(bias_op.out_tensors[0]) + net.graph_outputs[index] = opr.out_tensors[0] + opr.activation = bias_op.activation + index = net.all_oprs.index(bias_op) + del net.all_oprs[index] + del net._opr_ids[index] + + +@_register_tranformation_rule(TransformerRule.FUSE_FOR_DECONV_BIAS) +def _fuse_for_deconv_bias(net: IRGraph): + for opr in net.all_oprs: + if ( + opr.name == "Deconv2d" + and len(net.find_out_oprs(opr)) == 1 + and net.find_out_oprs(opr)[0].name == "Add" + ): + bias_op = net.find_out_oprs(opr)[0] + if not ( + ( + bias_op.inp_tensors[1].np_data is not None + and len(bias_op.inp_tensors[1].np_data.reshape(-1)) + == opr.inp_tensors[1].shape[1] + ) + or ( + ( + bias_op.inp_tensors[0].np_data is not None + and len(bias_op.inp_tensors[0].np_data.reshape(-1)) + == opr.inp_tensors[1].shape[1] + ) + ) + ): + continue + bias_idx = 0 if bias_op.inp_tensors[0].np_data is not None else 1 + if len(opr.inp_tensors) == 3: + opr.inp_tensors.append(bias_op.inp_tensors[bias_idx]) + else: + opr.inp_tensors[2].np_data += bias_op.inp_tensors[ + bias_idx + ].np_data.reshape(-1) + if bias_op in opr.out_tensors[0].user_opr: + opr.out_tensors[0].user_opr.remove(bias_op) + bias_out_op = net.find_out_oprs(bias_op) + if len(bias_out_op) > 0: + for op in bias_out_op: + op.inp_tensors[0] = opr.out_tensors[0] + opr.out_tensors[0].user_opr.append(op) + else: + # last op of the graph + assert bias_op.out_tensors[0] in net.graph_outputs + index = net.graph_outputs.index(bias_op.out_tensors[0]) + net.graph_outputs[index] = opr.out_tensors[0] + opr.activation = bias_op.activation + index = net.all_oprs.index(bias_op) + del net.all_oprs[index] + del net._opr_ids[index] + + +@_register_tranformation_rule(TransformerRule.EXPAND_MUL_ADD3) +def _expand_mul_add3(net: IRGraph): + + for op in net.all_oprs: + if not isinstance(op, FuseMulAdd3Opr): + continue + + last_op = net.find_inp_oprs(op) + assert len(last_op) == 1 + mul_out_tensor = IRTensor( + name=op.inp_tensors[0].name + "_mul_out", + shape=op.inp_tensors[0].shape, + dtype=op.inp_tensors[0].dtype, + ) + new_tensor_id = max(net._tensor_ids) + 1 + net.add_tensor(new_tensor_id, mul_out_tensor) + + mul_op = MulOpr() + mul_out_tensor.owner_opr = mul_op + mul_op.inp_tensors = op.inp_tensors[:2] + for o in mul_op.inp_tensors: + index = o.user_opr.index(op) + o.user_opr[index] = mul_op + mul_op.out_tensors = [mul_out_tensor] + + add_op = AddOpr() + add_op.inp_tensors = [mul_out_tensor, op.inp_tensors[2]] + mul_out_tensor.user_opr.append(add_op) + add_op.out_tensors = op.out_tensors + + index = net._opr_ids.index(id(op)) + net.delete_ops(index) + net.add_op(mul_op, index) + net.add_op(add_op, index + 1) + + +@_register_tranformation_rule(TransformerRule.REMOVE_RESHAPE_REALTED_OP) +def _remove_reshape_tensors(net: IRGraph): + for opr in net.all_oprs: + if isinstance(opr, ReshapeOpr) and len(opr.inp_tensors) > 1: + opr.inp_tensors = opr.inp_tensors[:1] + + +@_register_tranformation_rule(TransformerRule.REMOVE_DROPOUT) +def _remove_dropout(net: IRGraph): + for opr in net.all_oprs: + for idx, inp in enumerate(opr.inp_tensors): + owner_opr = inp.owner_opr + if isinstance(owner_opr, DropoutOpr) and owner_opr.drop_prob == 0: + opr.inp_tensors[idx] = owner_opr.inp_tensors[0] + + for idx, out in enumerate(net.graph_outputs): + owner_opr = out.owner_opr + if isinstance(owner_opr, DropoutOpr) and owner_opr.drop_prob == 0: + net.graph_outputs[idx] = owner_opr.inp_tensors[0] + + +visited_tensor = set() # type: set + + +def _dfs_recursive(op_set, tensor): + owner_opr = tensor.owner_opr + op_set.add(owner_opr) + if tensor in visited_tensor: + return + visited_tensor.add(tensor) + if isinstance(owner_opr, IRGraph) or owner_opr is None: + return + for tt in owner_opr.inp_tensors: + _dfs_recursive(op_set, tt) + + +@_register_tranformation_rule(TransformerRule.REMOVE_UNRELATED_IROP) +def _remove_unrelated_op(net: IRGraph): + match_sets = set() # type: Set[OpBase] + for out_tensor in net.graph_outputs: + _dfs_recursive(match_sets, out_tensor) + remove_idx = [] + for opr in net.all_oprs: + if opr not in match_sets: + index = net._opr_ids.index(id(opr)) + remove_idx.append(index) + for i in remove_idx[::-1]: + net.delete_ops(i) + + +@_register_tranformation_rule(TransformerRule.ADD_FAKE_HSIGMOID_OUT) +def _add_fake_hsigmoid_tensor(net: IRGraph): + for opr in net.all_oprs: + if isinstance(opr, (HardSwishOpr, HardSigmoidOpr)): + add_3_out_tensor = IRTensor( + opr.inp_tensors[0].name + "_fake_add3_out", + opr.inp_tensors[0].shape, + opr.inp_tensors[0].dtype, + ) + opr.add_inp_tensors(add_3_out_tensor) + relu6_out_tensor = IRTensor( + opr.inp_tensors[0].name + "_relu6_out", + opr.inp_tensors[0].shape, + opr.inp_tensors[0].dtype, + ) + opr.add_inp_tensors(relu6_out_tensor) + if isinstance(opr, HardSwishOpr): + div6_out_tensor = IRTensor( + opr.inp_tensors[0].name + "_div_out", + opr.inp_tensors[0].shape, + opr.inp_tensors[0].dtype, + ) + opr.add_inp_tensors(div6_out_tensor) diff --git a/mgeconvert/caffe_converter/__init__.py b/mgeconvert/converters/__init__.py similarity index 83% rename from mgeconvert/caffe_converter/__init__.py rename to mgeconvert/converters/__init__.py index 0ccdd96..93515f3 100644 --- a/mgeconvert/caffe_converter/__init__.py +++ b/mgeconvert/converters/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,4 +5,3 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .caffe_converter import convert_to_caffe diff --git a/mgeconvert/converters/mge_to_caffe.py b/mgeconvert/converters/mge_to_caffe.py new file mode 100644 index 0000000..8f736ea --- /dev/null +++ b/mgeconvert/converters/mge_to_caffe.py @@ -0,0 +1,36 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ..backend.ir_to_caffe import CaffeConverter +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.mge_to_ir import MGE_FrontEnd + + +def mge_to_caffe( + mge_fpath, + prototxt="out.prototxt", + caffemodel="out.caffemodel", + outspec=None, + use_empty_blobs=False, +): + assert isinstance(mge_fpath, str), "mge_fpath must be string" + irgraph = MGE_FrontEnd(mge_fpath, outspec=outspec).resolve() + + transformer_options = [ + TransformerRule.EXPAND_MUL_ADD3, + TransformerRule.FUSE_FOR_LEAKY_RELU, + ] + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + + converter = CaffeConverter(transformed_irgraph, use_empty_blobs) + converter.convert() + + assert isinstance(prototxt, str) and isinstance( + caffemodel, str + ), "'prototxt' and 'caffemodel' must be string" + converter.dump(prototxt, caffemodel) diff --git a/mgeconvert/converters/mge_to_onnx.py b/mgeconvert/converters/mge_to_onnx.py new file mode 100644 index 0000000..8d30374 --- /dev/null +++ b/mgeconvert/converters/mge_to_onnx.py @@ -0,0 +1,49 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ..backend.ir_to_onnx import OnnxConverter +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.mge_to_ir import MGE_FrontEnd + + +def remove_initializer_from_input(model): + if model.ir_version < 4: + print( + "Model with ir_version below 4 requires to include initilizer in graph input" + ) + return model + + inputs = model.graph.input + name_to_input = {} + for input in inputs: + name_to_input[input.name] = input + + for initializer in model.graph.initializer: + if initializer.name in name_to_input: + inputs.remove(name_to_input[initializer.name]) + + return model + + +def mge_to_onnx( + mge_fpath, output="out.onnx", *, graph_name="graph", opset=8, outspec=None +): + assert isinstance(mge_fpath, str), "mge_fpath must be string" + irgraph = MGE_FrontEnd(mge_fpath, outspec=outspec).resolve() + transformer_options = [ + TransformerRule.FUSE_SOFTMAX, + TransformerRule.EXPAND_MUL_ADD3, + ] + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + converter = OnnxConverter(transformed_irgraph, opset, graph_name) + model = converter.convert() + model = remove_initializer_from_input(model) + + assert isinstance(output, str), "onnx_fpath must be string" + with open(output, "wb") as fout: + fout.write(model.SerializeToString()) diff --git a/mgeconvert/converters/mge_to_tflite.py b/mgeconvert/converters/mge_to_tflite.py new file mode 100644 index 0000000..7c58ae3 --- /dev/null +++ b/mgeconvert/converters/mge_to_tflite.py @@ -0,0 +1,58 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ..backend.ir_to_tflite import TFLiteConverter, set_platform +from ..converter_ir.ir_quantizer import IRQuantizer +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.mge_to_ir import MGE_FrontEnd + + +def mge_to_tflite( + mge_fpath, + output="out.tflite", + *, + graph_name="graph", + mtk=False, + disable_nhwc=False, +): + assert isinstance(mge_fpath, str), "mge_fpath must be string" + irgraph = MGE_FrontEnd(mge_fpath).resolve() + + transformer_options = [ + TransformerRule.REDUCE_AXIS_AS_INPUT, + TransformerRule.PADDING_FOR_CONV_AND_POOLING, + TransformerRule.CONV_ADD_ZERO_BIAS, + TransformerRule.DECONV_SHAPE_AS_INPUT, + TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT, + TransformerRule.RESHAPE_BIAS_TO_1DIM, + TransformerRule.FUSE_ACTIVATION, + TransformerRule.SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE, + TransformerRule.RESIZE_PARAMS_AS_INPUT, + TransformerRule.TRANSPOSE_PATTERN_AS_INPUT, + TransformerRule.REMOVE_RESHAPE_INPUT, + TransformerRule.FUSE_SOFTMAX, + TransformerRule.EXPAND_MUL_ADD3, + TransformerRule.FUSE_FOR_CONV_BIAS, + TransformerRule.FUSE_FOR_LEAKY_RELU, + ] + if mtk: + # MTK devices only support batch_size 1 + set_platform("mtk") + transformer_options.append(TransformerRule.DECONV_ADD_ZERO_BIAS,) + transformer_options.append(TransformerRule.FUSE_FOR_DECONV_BIAS,) + + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + + quantizer = IRQuantizer(require_quantize=False) + + converter = TFLiteConverter(transformed_irgraph, graph_name, quantizer=quantizer) + model = converter.convert(disable_nhwc=disable_nhwc) + + assert isinstance(output, str), "tflite_fpath must be string" + with open(output, "wb") as fout: + fout.write(model) diff --git a/mgeconvert/converters/tm_to_caffe.py b/mgeconvert/converters/tm_to_caffe.py new file mode 100644 index 0000000..e0a3a1e --- /dev/null +++ b/mgeconvert/converters/tm_to_caffe.py @@ -0,0 +1,57 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +import megengine as mge +from megengine.traced_module import TracedModule + +from ..backend.ir_to_caffe.caffe_converter import CaffeConverter +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.tm_to_ir import TM_FrontEnd + + +def tracedmodule_to_caffe( + traced_module, + prototxt="out.prototxt", + caffemodel="out.caffemodel", + use_empty_blobs=False, +): + """ + Convert megengine model to Caffe, + and save caffe model to `prototxt` and `caffemodel`. + + :param mge_fpath: the file path of megengine model. + :type mge_fpath: str + :param prototxt: the filename used for saved model definition. + :type prototxt: str + :param caffemodel: the filename used for saved model weights. + :type caffemodel: str + """ + if isinstance(traced_module, str): + traced_module = mge.load(traced_module) + assert isinstance( + traced_module, TracedModule + ), "Input should be a traced module or a path of traced module." + + irgraph = TM_FrontEnd(traced_module).resolve() + transformer_options = [ + TransformerRule.REMOVE_DROPOUT, + TransformerRule.REMOVE_RESHAPE_REALTED_OP, + TransformerRule.REMOVE_UNRELATED_IROP, + TransformerRule.ADD_FAKE_HSIGMOID_OUT, + ] + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + converter = CaffeConverter(transformed_irgraph, use_empty_blobs) + converter.convert() + + assert isinstance(prototxt, str) and isinstance( + caffemodel, str + ), "'prototxt' and 'caffemodel' must be string" + converter.dump(prototxt, caffemodel) diff --git a/mgeconvert/converters/tm_to_onnx.py b/mgeconvert/converters/tm_to_onnx.py new file mode 100644 index 0000000..e423d03 --- /dev/null +++ b/mgeconvert/converters/tm_to_onnx.py @@ -0,0 +1,54 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +import megengine as mge +from megengine.traced_module import TracedModule + +from ..backend.ir_to_onnx.onnx_converter import OnnxConverter +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.tm_to_ir import TM_FrontEnd + + +def tracedmodule_to_onnx( + traced_module, output="out.onnx", *, graph_name="graph", opset=8 +): + """ + Convert megengine model to ONNX, + and save the ONNX model to file `output`. + + :param mge_fpath: the file path of megengine model. + :type fpath: str + :param output: the filename used for the saved model. + :type output: str + :param graph_name: the name of the ONNX graph. + :type graph_name: str + :param opset: opset version of ONNX model. + :type opset: int + """ + if isinstance(traced_module, str): + traced_module = mge.load(traced_module) + assert isinstance( + traced_module, TracedModule + ), "Input should be a traced module or a path of traced module." + + irgraph = TM_FrontEnd(traced_module).resolve() + transformer_options = [ + TransformerRule.REMOVE_RESHAPE_REALTED_OP, + TransformerRule.REMOVE_UNRELATED_IROP, + ] + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + + converter = OnnxConverter(transformed_irgraph, opset, graph_name) + model = converter.convert() + + assert isinstance(output, str), "onnx_fpath must be string" + with open(output, "wb") as fout: + fout.write(model.SerializeToString()) diff --git a/mgeconvert/converters/tm_to_tflite.py b/mgeconvert/converters/tm_to_tflite.py new file mode 100644 index 0000000..e8ee03f --- /dev/null +++ b/mgeconvert/converters/tm_to_tflite.py @@ -0,0 +1,110 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member +from typing import List, Sequence, Union + +import megengine as mge +from megengine.core.tensor import dtype +from megengine.quantization.utils import create_qparams +from megengine.traced_module import TracedModule + +from ..backend.ir_to_tflite import TFLiteConverter, set_platform +from ..converter_ir.ir_quantizer import IRQuantizer +from ..converter_ir.ir_transform import IRTransform, TransformerRule +from ..frontend.tm_to_ir import TM_FrontEnd + + +def tracedmodule_to_tflite( + traced_module, + output="out.tflite", + *, + input_data_type: str = None, + input_scales: Union[float, List[float]] = None, + input_zero_points: Union[int, List[int]] = None, + require_quantize=False, + param_fake_quant=False, + quantize_file_path="quant_params.json", + graph_name="graph", + mtk=False, +): + """ + Convert traced model to TFLite, + and save the TFLite model to file `output`. + + :param traced_module: a traced module or the file path of a traced module. + :param output: the filename used for the saved model. + :param data_type: data type of input + + :param graph_name: the name of the TFLite graph. + :param mtk: if this TFLite will be run on mtk. + :type mtk: bool + """ + if isinstance(traced_module, str): + traced_module = mge.load(traced_module) + assert isinstance( + traced_module, TracedModule + ), "Input should be a traced module or a path of traced module." + + if input_data_type is not None: + for i in range(len(traced_module.graph.inputs[1:])): + if traced_module.graph.inputs[i + 1].qparams is None: + traced_module.graph.inputs[i + 1].qparams = create_qparams() + traced_module.graph.inputs[ + i + 1 + ].qparams.dtype_meta = dtype._builtin_quant_dtypes[input_data_type] + if input_scales is not None: + if not isinstance(input_scales, Sequence): + scales = (input_scales,) + for i in range(len(traced_module.graph.inputs[1:])): + scale = scales[i] if i < len(scales) else scales[-1] + traced_module.graph.inputs[i + 1].qparams.scale = mge.tensor(float(scale)) + if input_zero_points is not None: + if not isinstance(input_zero_points, Sequence): + zero_points = (input_zero_points,) + for i in range(len(traced_module.graph.inputs[1:])): + zero_point = zero_points[i] if i < len(zero_points) else zero_points[-1] + traced_module.graph.inputs[i + 1].qparams.zero_point = mge.tensor( + int(zero_point) + ) + + irgraph = TM_FrontEnd(traced_module).resolve() + + transformer_options = [ + TransformerRule.REDUCE_AXIS_AS_INPUT, + TransformerRule.PADDING_FOR_CONV_AND_POOLING, + TransformerRule.CONV_ADD_ZERO_BIAS, + TransformerRule.DECONV_SHAPE_AS_INPUT, + TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT, + TransformerRule.RESHAPE_BIAS_TO_1DIM, + TransformerRule.FUSE_ACTIVATION, + TransformerRule.SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE, + TransformerRule.RESIZE_PARAMS_AS_INPUT, + TransformerRule.TRANSPOSE_PATTERN_AS_INPUT, + ] + if mtk: + # MTK devices only support batch_size 1 + set_platform("mtk") + transformer_options.append(TransformerRule.DECONV_ADD_ZERO_BIAS,) + + transformer = IRTransform(transformer_options) + transformed_irgraph = transformer.transform(irgraph) + + quantizer = IRQuantizer( + require_quantize=require_quantize, param_fake_quant=param_fake_quant + ) + + if not require_quantize: + quantizer.save_quantize_params(transformed_irgraph, path=quantize_file_path) + + converter = TFLiteConverter(transformed_irgraph, graph_name, quantizer=quantizer) + model = converter.convert() + + assert isinstance(output, str), "tflite_fpath must be string" + with open(output, "wb") as fout: + fout.write(model) diff --git a/mgeconvert/cambricon_converter/__init__.py b/mgeconvert/frontend/__init__.py similarity index 81% rename from mgeconvert/cambricon_converter/__init__.py rename to mgeconvert/frontend/__init__.py index 8c88c0e..93515f3 100644 --- a/mgeconvert/cambricon_converter/__init__.py +++ b/mgeconvert/frontend/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,4 +5,3 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .cambricon_converter import convert_to_cambricon diff --git a/mgeconvert/frontend/mge_to_ir/__init__.py b/mgeconvert/frontend/mge_to_ir/__init__.py new file mode 100644 index 0000000..a4db943 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/__init__.py @@ -0,0 +1,8 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from .mge_frontend import MGE_FrontEnd diff --git a/mgeconvert/frontend/mge_to_ir/mge_frontend.py b/mgeconvert/frontend/mge_to_ir/mge_frontend.py new file mode 100644 index 0000000..3c1a53b --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/mge_frontend.py @@ -0,0 +1,65 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ...converter_ir.ir_graph import IRGraph +from .mge_utils import ( + get_dep_vars, + get_opr_type, + get_oprs_seq, + load_comp_graph_from_file, +) +from .op_generators import MGE2OP +from .symbolvar_resolver import SymbolVarResolver + + +class MGE_FrontEnd: + def __init__(self, model_path, outspec=None, prune_reshape=True): + _, outputs = load_comp_graph_from_file(model_path) + if outspec is not None: + output_spec = outspec.copy() + all_vars = get_dep_vars(outputs) + outputs + new_outputs = {} + for i in all_vars: + if i.name in output_spec: + new_outputs[i.name] = i + output_spec.remove(i.name) + assert len(output_spec) == 0, "Can not find {} in this model".format( + output_spec + ) + outputs = [new_outputs[i] for i in outspec] + self.ori_all_oprs = get_oprs_seq(outputs, prune_reshape=prune_reshape) + self._orig_inputs = [] + self._orig_outputs = outputs + self.opr_maps = {} # {mge_opr.id: id(ir_opr)} + + self.irgraph = IRGraph() + self.resolver = SymbolVarResolver(self.irgraph) + + def resolve(self): + for mge_opr in self.ori_all_oprs: + if get_opr_type(mge_opr) == "Host2DeviceCopy": + # add graph inputs + for inp_var in mge_opr.outputs: + self.irgraph.add_net_inputs(self.resolver.get_ir_tensor(inp_var)) + continue + self.add_opr(mge_opr) + + # add graph outputs + for out_var in self._orig_outputs: + self.irgraph.add_net_outputs(self.resolver.get_ir_tensor(out_var)) + + return self.irgraph + + def add_opr(self, mge_opr): + assert mge_opr.id not in self.opr_maps.keys() + + op_gen_cls = MGE2OP.get(get_opr_type(mge_opr), None) + assert op_gen_cls, "OP {} is not supported".format(get_opr_type(mge_opr)) + ir_op = op_gen_cls(mge_opr, self.irgraph).get_opr() + + self.irgraph.add_op(ir_op) + self.opr_maps[mge_opr.id] = id(ir_op) diff --git a/mgeconvert/mge_context/mge_utils.py b/mgeconvert/frontend/mge_to_ir/mge_utils.py similarity index 99% rename from mgeconvert/mge_context/mge_utils.py rename to mgeconvert/frontend/mge_to_ir/mge_utils.py index 54644ac..cdf07e7 100644 --- a/mgeconvert/mge_context/mge_utils.py +++ b/mgeconvert/frontend/mge_to_ir/mge_utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -18,8 +17,8 @@ import megengine._internal as mgb from megengine._internal import cgtools else: - import megengine.core.tensor.megbrain_graph as G import megengine.core._imperative_rt as rt + import megengine.core.tensor.megbrain_graph as G import megengine.utils.comp_graph_tools as cgtools if mge_version <= "1.1.0": diff --git a/mgeconvert/cambricon_converter/lib/cnlib/__init__.py b/mgeconvert/frontend/mge_to_ir/op_generators/__init__.py similarity index 76% rename from mgeconvert/cambricon_converter/lib/cnlib/__init__.py rename to mgeconvert/frontend/mge_to_ir/op_generators/__init__.py index ab17733..7414d55 100644 --- a/mgeconvert/cambricon_converter/lib/cnlib/__init__.py +++ b/mgeconvert/frontend/mge_to_ir/op_generators/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,4 +5,7 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from . import cambriconLib # type: ignore[attr-defined] # pylint: disable=import-self +from .base import MGE2OP +from .convolution import * +from .elemwise import GenElemwiseOpr +from .tensor import * diff --git a/mgeconvert/frontend/mge_to_ir/op_generators/base.py b/mgeconvert/frontend/mge_to_ir/op_generators/base.py new file mode 100644 index 0000000..bd23ac4 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/op_generators/base.py @@ -0,0 +1,62 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +import json +from abc import ABC + +from mgeconvert.converter_ir.ir_op import OpBase + +from ..mge_utils import get_mge_version, get_opr_type +from ..symbolvar_resolver import SymbolVarResolver + +mge_version = get_mge_version() + +MGE2OP = {} + + +def _register_op(*ops): + def callback(impl): + for op in ops: + MGE2OP[op] = impl + return impl + + return callback + + +class OpGenBase(ABC): + def __init__(self, mge_opr, irgraph) -> None: + self.mge_opr = mge_opr + self.irgraph = irgraph + self.name = mge_opr.name + self.name = self.name.replace(":", "_") + self.name = self.name.replace(".", "_") + self.name = self.name.replace(",", "_") + + self.id = mge_opr.id + try: + self.type = get_opr_type(mge_opr) + except AssertionError: + self.type = None + self.flag = None + self.params = ( + mge_opr.params if mge_version <= "0.6.0" else json.loads(mge_opr.params) + ) + self.activation = "IDENTITY" + + self.resolver = SymbolVarResolver(irgraph) + self.op = OpBase() + + def get_opr(self): + return self.op + + def add_tensors(self, mge_opr): + # set inp var + for x in mge_opr.inputs: + self.op.add_inp_tensors(self.resolver.get_ir_tensor(x, user_opr=self.op)) + # set out var + for x in mge_opr.outputs: + self.op.add_out_tensors(self.resolver.get_ir_tensor(x, owner_opr=self.op)) diff --git a/mgeconvert/frontend/mge_to_ir/op_generators/convolution.py b/mgeconvert/frontend/mge_to_ir/op_generators/convolution.py new file mode 100644 index 0000000..46e1365 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/op_generators/convolution.py @@ -0,0 +1,174 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ....converter_ir.ir_op import ( + AvgPool2dOpr, + BatchNormalizationOpr, + Conv2dOpr, + ConvolutionBackwardFilterOpr, + Deconv2dOpr, + MaxPool2dOpr, +) +from ....converter_ir.ir_tensor import AxisOrder +from ..mge_utils import get_shape +from .base import OpGenBase, _register_op + + +@_register_op("ConvolutionForward") +class GenConv2dOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.kernel_shape = get_shape(mge_opr.inputs[1]) + + self.stride = [self.params["stride_h"], self.params["stride_w"]] + self.padding = [self.params["pad_h"], self.params["pad_w"]] + self.dilation = [self.params["dilate_h"], self.params["dilate_w"]] + self.groups = self.kernel_shape[0] if len(self.kernel_shape) == 5 else 1 + + self.data_format = self.params["format"] + assert self.data_format == "NCHW", "do not support this {} format".format( + self.data_format + ) + assert self.params["compute_mode"].lower() == "default" + assert self.params["mode"].lower() == "cross_correlation" + + self.op = Conv2dOpr(self.stride, self.padding, self.dilation, self.groups) + self.add_tensors(mge_opr) + + def add_tensors(self, mge_opr): + self.op.add_inp_tensors( + self.resolver.get_ir_tensor(mge_opr.inputs[0], user_opr=self.op) + ) + self.op.add_inp_tensors( + self.resolver.get_ir_tensor( + mge_opr.inputs[1], user_opr=self.op, axis_order=AxisOrder.OIHW + ) + ) + if len(mge_opr.inputs) > 2: + self.op.add_inp_tensors( + self.resolver.get_ir_tensor(mge_opr.inputs[2], user_opr=self.op) + ) + + for x in mge_opr.outputs: + self.op.add_out_tensors(self.resolver.get_ir_tensor(x, owner_opr=self.op)) + + +@_register_op("ConvBiasForward") +class GenConvBiasForwardOpr(GenConv2dOpr): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op.activation = self.params["nonlineMode"] + + +@_register_op("ConvolutionBackwardData") +class GenDeconv2dOpr(OpGenBase): + def __init__(self, mge_opr, irgraph) -> None: + super().__init__(mge_opr, irgraph) + self.kernel_shape = get_shape(mge_opr.inputs[0]) + + self.stride = [self.params["stride_h"], self.params["stride_w"]] + self.padding = [self.params["pad_h"], self.params["pad_w"]] + self.dilation = [self.params["dilate_h"], self.params["dilate_w"]] + self.groups = self.kernel_shape[0] if len(self.kernel_shape) == 5 else 1 + + self.sparse = self.params["sparse"] + self.data_format = self.params["format"] + assert self.data_format == "NCHW", "do not support this {} format".format( + self.data_format + ) + assert self.params["compute_mode"].lower() == "default" + assert self.params["mode"].lower() == "cross_correlation" + + self.op = Deconv2dOpr(self.stride, self.padding, self.dilation, self.groups) + self.add_tensors(mge_opr) + + def add_tensors(self, mge_opr): + self.op.add_inp_tensors( + self.resolver.get_ir_tensor(mge_opr.inputs[1], user_opr=self.op) + ) + self.op.add_inp_tensors( + self.resolver.get_ir_tensor( + mge_opr.inputs[0], user_opr=self.op, axis_order=AxisOrder.IOHW + ) + ) + if len(mge_opr.inputs) > 2: + self.op.add_inp_tensors( + self.resolver.get_ir_tensor(mge_opr.inputs[2], user_opr=self.op) + ) + + for x in mge_opr.outputs: + self.op.add_out_tensors(self.resolver.get_ir_tensor(x, owner_opr=self.op)) + + +mode_map = {"MAX": MaxPool2dOpr, "AVERAGE_COUNT_EXCLUDE_PADDING": AvgPool2dOpr} + + +@_register_op("PoolingForward") +class GenPool2dOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.data_format = self.params["format"] + self.mode = self.params["mode"] + self.stride = [self.params["stride_h"], self.params["stride_w"]] + self.padding = [self.params["pad_h"], self.params["pad_w"]] + self.kernel_shape = [self.params["window_h"], self.params["window_w"]] + + self.op = mode_map[self.mode](self.kernel_shape, self.stride, self.padding) + self.add_tensors(mge_opr) + + +@_register_op("ConvolutionBackwardFilter") +class GenConvolutionBackwardFilterOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + assert self.params["format"] == "NCHW", "do not support this {}".format( + self.params["format"] + ) + src, grad_out, weight = mge_opr.inputs + + self.ni, self.ci, self.hi, self.wi = src.shape + self.no, self.co, self.ho, self.wo = grad_out.shape + assert self.ni == self.no + + if len(weight.shape) == 5: + self.group = weight.shape[0] + self.kernel_shape = [weight.shape[3], weight.shape[4]] + else: + self.group = 1 + self.kernel_shape = [weight.shape[2], weight.shape[3]] + + self.stride = [self.params["stride_h"], self.params["stride_w"]] + self.padding = [self.params["pad_h"], self.params["pad_w"]] + self.dilation = [self.params["dilate_h"], self.params["dilate_w"]] + + self.op = ConvolutionBackwardFilterOpr( + self.stride, + self.padding, + self.dilation, + self.group, + self.kernel_shape, + src.shape, + grad_out.shape, + ) + self.add_tensors(mge_opr) + + +@_register_op("BatchNormForward") +class GenBatchNormalizationOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.output_idx = -1 + self.op = BatchNormalizationOpr(output_idx=self.output_idx) + self.add_tensors(mge_opr) + + def add_tensors(self, mge_opr): + # set inp var + for x in mge_opr.inputs: + self.op.add_inp_tensors(self.resolver.get_ir_tensor(x, user_opr=self.op)) + # set out var + for x in mge_opr.outputs: + self.op.add_out_tensors(self.resolver.get_ir_tensor(x, owner_opr=self.op)) diff --git a/mgeconvert/frontend/mge_to_ir/op_generators/elemwise.py b/mgeconvert/frontend/mge_to_ir/op_generators/elemwise.py new file mode 100644 index 0000000..a1467f6 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/op_generators/elemwise.py @@ -0,0 +1,71 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ....converter_ir.ir_op import ( + AbsOpr, + AddOpr, + CeilOpr, + ExpOpr, + FloorOpr, + FuseMulAdd3Opr, + LogOpr, + MaxOpr, + MinOpr, + MulOpr, + PowOpr, + ReluOpr, + SigmoidOpr, + SubOpr, + TanHOpr, + TrueDivOpr, +) +from .base import OpGenBase, _register_op + + +class GenFuseMulAdd3Oprs(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = FuseMulAdd3Opr() + self.add_tensors(mge_opr) + + +mode_opr_map = { + "add": AddOpr, + "fuse_add_relu": AddOpr, + "sigmoid": SigmoidOpr, + "mul": MulOpr, + "abs": AbsOpr, + "ceil": CeilOpr, + "exp": ExpOpr, + "floor": FloorOpr, + "log": LogOpr, + "max": MaxOpr, + "min": MinOpr, + "pow": PowOpr, + "relu": ReluOpr, + "sub": SubOpr, + "tanh": TanHOpr, + "true_div": TrueDivOpr, + "fuse_mul_add3": GenFuseMulAdd3Oprs, +} + + +@_register_op("Elemwise") +class GenElemwiseOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + try: + self.mode = self.params["mode"] + except RuntimeError: + self.mode = "NONE" + if self.mode.lower() in ["fuse_mul_add3"]: + self.op = mode_opr_map[self.mode.lower()](mge_opr, irgraph).get_opr() + else: + self.op = mode_opr_map[self.mode.lower()]() + if "RELU" in self.mode: + self.op.activation = "RELU" + self.add_tensors(mge_opr) diff --git a/mgeconvert/frontend/mge_to_ir/op_generators/tensor.py b/mgeconvert/frontend/mge_to_ir/op_generators/tensor.py new file mode 100644 index 0000000..8582795 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/op_generators/tensor.py @@ -0,0 +1,221 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ....converter_ir.ir_op import ( + AxisAddRemoveOpr, + BroadcastOpr, + ConcatOpr, + ConstantOpr, + GetSubTensorOpr, + GetVarShapeOpr, + IdentityOpr, + MatMulOpr, + MultipleDeviceTensorHolderOpr, + ReduceOpr, + ReshapeOpr, + ResizeOpr, + SharedDeviceTensorOpr, + TransposeOpr, + TypeCvtOpr, + VolatileSharedDeviceTensorOpr, +) +from ..mge_utils import get_shape, get_symvar_value +from .base import OpGenBase, _register_op + + +@_register_op("ImmutableTensor") +class GenImmutableTensorOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = ConstantOpr() + self.add_tensors(mge_opr) + + +@_register_op("MultipleDeviceTensorHolder") +class GenMultipleDeviceTensorHolderOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = MultipleDeviceTensorHolderOpr() + self.add_tensors(mge_opr) + + +@_register_op("SharedDeviceTensorOpr") +class GenSharedDeviceTensorOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = SharedDeviceTensorOpr() + self.add_tensors(mge_opr) + + +@_register_op("VolatileSharedDeviceTensor") +class GenVolatileSharedDeviceTensorOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = VolatileSharedDeviceTensorOpr() + self.add_tensors(mge_opr) + + +@_register_op("Identity") +class GenIdentityOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = IdentityOpr() + self.add_tensors(mge_opr) + + +@_register_op("GetVarShape") +class GenGetVarShapeOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = GetVarShapeOpr() + self.add_tensors(mge_opr) + self.op.out_tensors[0]._var = mge_opr.outputs[0] + + +@_register_op("Broadcast") +class GenBroadcastOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.op = BroadcastOpr() + self.add_tensors(mge_opr) + for i in range(len(self.op.out_tensors)): + self.op.out_tensors[i]._var = mge_opr.outputs[i] + + +@_register_op("Concat") +class GenConcatOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.axis = self.params["axis"] + self.op = ConcatOpr(self.axis) + self.add_tensors(mge_opr) + + +@_register_op("Reshape") +class GenReshapeOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.output_shape = get_shape(mge_opr.outputs[0]) + + self.op = ReshapeOpr(self.output_shape) + self.add_tensors(mge_opr) + + +@_register_op("AxisAddRemove") +class GenAxisAddRemoveOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.desc = self.params["desc"] # 0:add_axis, 1:remove_axis + self.nr_desc = self.params["nr_desc"] + self.output_shape = get_shape(mge_opr.outputs[0]) + self.op = AxisAddRemoveOpr(self.output_shape, self.desc) + self.add_tensors(mge_opr) + + +@_register_op("Reduce") +class GenReduceOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.axis = self.params["axis"] + self.mode = self.params["mode"] + self.op = ReduceOpr(self.axis, self.mode, True) + self.add_tensors(mge_opr) + for i in range(len(self.op.out_tensors)): + self.op.out_tensors[i]._var = mge_opr.outputs[i] + + +@_register_op("MatrixMul") +class GenMatrixMulOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.format = self.params["format"] + self.transposeB = self.params["transposeB"] + self.transposeA = self.params["transposeA"] + self.compute_mode = self.params["compute_mode"] + + self.op = MatMulOpr( + transpose_a=self.transposeA, + transpose_b=self.transposeB, + compute_mode=self.compute_mode, + format=self.format, + ) + self.add_tensors(mge_opr) + + +@_register_op("Dimshuffle") +class GenDimshuffleOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.pattern = self.params["pattern"] + + self.op = TransposeOpr(self.pattern) + self.add_tensors(mge_opr) + + +@_register_op("Subtensor") +class GenSubtensorOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.has_step = [] + self.has_begin = [] + self.has_end = [] + self.has_idx = [] + self.axis = [] + for param in self.params: + self.has_step.append(param["step"]) + self.has_begin.append(param["begin"]) + self.has_end.append(param["end"]) + self.has_idx.append(param["idx"]) + self.axis.append(param["axis"]) + + begin_param = [] + end_param = [] + step_param = [] + squeeze_axis = [] + slice_param = [get_symvar_value(v)[0] for v in mge_opr.inputs[1:]] + for i in range(len(self.has_begin)): + if self.has_idx[i] == 1: + begin_idx = slice_param.pop(0) + end_idx = begin_idx + 1 + begin_param.append(begin_idx) + end_param.append(end_idx) + step_param.append(1) + squeeze_axis.append(self.axis[i]) + else: + if self.has_begin[i]: + begin_param.append(slice_param.pop(0)) + else: + begin_param.append(0) + if self.has_end[i]: + end_param.append(slice_param.pop(0)) + else: + end_param.append(2147483647) + step_param.append(1 if self.has_step[i] == 0 else slice_param.pop(0)) + + self.op = GetSubTensorOpr( + self.axis, begin_param, end_param, step_param, squeeze_axis=squeeze_axis + ) + self.add_tensors(mge_opr) + + +@_register_op("TypeCvt") +class GenTypeCvtOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + + out_dtype = mge_opr.outputs[0].dtype + self.op = TypeCvtOpr(out_dtype) + self.add_tensors(mge_opr) + + +@_register_op("ResizeForward") +class GenResizeOpr(OpGenBase): + def __init__(self, mge_opr, irgraph): + super().__init__(mge_opr, irgraph) + self.shape_param = get_symvar_value(mge_opr.inputs[1]) + self.op = ResizeOpr(self.shape_param) + self.add_tensors(mge_opr) diff --git a/mgeconvert/frontend/mge_to_ir/symbolvar_resolver.py b/mgeconvert/frontend/mge_to_ir/symbolvar_resolver.py new file mode 100644 index 0000000..951b056 --- /dev/null +++ b/mgeconvert/frontend/mge_to_ir/symbolvar_resolver.py @@ -0,0 +1,64 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ...converter_ir.ir_graph import IRGraph +from ...converter_ir.ir_tensor import AxisOrder, IRTensor +from .mge_utils import get_dep_vars, get_shape, get_symvar_value + + +class SymbolVarResolver: + def __init__(self, irgraph: IRGraph, axis_order=AxisOrder.NCHW) -> None: + self.irgraph = irgraph + self.axis_order = axis_order + + def resolve(self, sym_var, owner_opr=None, axis_order=AxisOrder.NCHW): + name = sym_var.name + name = name.replace(":", "_") + name = name.replace(".", "_") + name = name.replace(",", "_") + + try: + dtype = sym_var.dtype + except: # pylint: disable=bare-except + dtype = None + + np_data = None + try: + if len(get_dep_vars(sym_var, "Host2DeviceCopy")) == 0: + np_data = get_symvar_value(sym_var) + except: # pylint: disable=bare-except + np_data = None + + try: + scale = dtype.metadata["mgb_dtype"]["scale"] + except: # pylint: disable=bare-except + scale = None + + try: + zero_point = dtype.metadata["mgb_dtype"]["zero_point"] + except: # pylint: disable=bare-except + zero_point = None + + return IRTensor( + name=name, + shape=get_shape(sym_var), + dtype=dtype, + scale=scale, + zero_point=zero_point, + np_data=np_data, + owner_opr=owner_opr, + axis=axis_order, + ) + + def get_ir_tensor( + self, sym_var, owner_opr=None, user_opr=None, axis_order=AxisOrder.NCHW, + ): + ir_tensor = self.resolve(sym_var, owner_opr=owner_opr, axis_order=axis_order) + ori_tensor = self.irgraph.get_tensor(sym_var.id, ir_tensor) + if user_opr is not None and user_opr not in ori_tensor.user_opr: + ori_tensor.add_user_opr(user_opr) + return ori_tensor diff --git a/mgeconvert/frontend/tm_to_ir/__init__.py b/mgeconvert/frontend/tm_to_ir/__init__.py new file mode 100644 index 0000000..9ac1074 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/__init__.py @@ -0,0 +1,8 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from .tm_frontend import TM_FrontEnd diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/__init__.py b/mgeconvert/frontend/tm_to_ir/op_generators/__init__.py new file mode 100644 index 0000000..88294f7 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/__init__.py @@ -0,0 +1,29 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +from .base import EXPR2OP +from .batchnorm import GenBatchNormalizationOpr +from .broadcast import GenBroadcastOpr +from .concat import GenConcatOpr, GenQConcatOpr +from .constant import ConstantOpr +from .conv2d import GenConv2dOpr, GenQConv2dOpr +from .deconv import GenDeconv2dOpr, GenQDeconv2dOpr +from .dropout import GenDropoutOpr +from .elemwise import * +from .flatten import GenFlattenOpr +from .getvarshape import GenGetVarShapeOpr +from .matmul import * +from .pooling import * +from .reduce import GenReduceOpr +from .reshape import GenRepeatOpr, GenReshapeOpr +from .resize import GenResizeOpr +from .softmax import GenSoftmaxOpr +from .squeeze import GenSqueezeOpr +from .subtensor import GenGetSubtensorOpr +from .transpose import GenTransposeOpr +from .typecvt import GenTypeCvtOpr diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/base.py b/mgeconvert/frontend/tm_to_ir/op_generators/base.py new file mode 100644 index 0000000..e4bf152 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/base.py @@ -0,0 +1,50 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from abc import ABC + +from mgeconvert.converter_ir.ir_op import OpBase + +from ..tm_tensor_resolver import TensorNodeResolver + +EXPR2OP = {} + + +def _register_op(*ops): + def callback(impl): + for op in ops: + EXPR2OP[op] = impl + return impl + + return callback + + +class OpGenBase(ABC): + def __init__(self, expr, irgraph) -> None: + self.expr = expr + self.irgraph = irgraph + self.resolver = TensorNodeResolver(self.irgraph) + self.op = OpBase() + + def get_opr(self): + return self.op + + def add_opr_out_tensors(self): + if len(self.op.inp_tensors) > 0: + is_qat = ( + hasattr(self.op.inp_tensors[0], "scale") + and self.op.inp_tensors[0].scale is not None + ) + else: + is_qat = False + for o in self.expr.outputs: + t = self.resolver.get_ir_tensor(o, owner_opr=self.op) + if is_qat: + t.scale = self.op.inp_tensors[0].scale + t.zero_point = self.op.inp_tensors[0].zero_point + t.q_dtype = self.op.inp_tensors[0].q_dtype + self.op.add_out_tensors(t) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/batchnorm.py b/mgeconvert/frontend/tm_to_ir/op_generators/batchnorm.py new file mode 100644 index 0000000..4ee610c --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/batchnorm.py @@ -0,0 +1,88 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import megengine.module as M +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import BatchNormalizationOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(M.BatchNorm1d, M.BatchNorm2d, F.batch_norm) +class GenBatchNormalizationOpr(OpGenBase): + def __init__(self, expr, ir_graph): + super().__init__(expr, ir_graph) + if isinstance(self.expr, CallMethod): + bn_module = expr.inputs[0].expr.value + state_dict = bn_module.state_dict() + self.running_mean = state_dict["running_mean"].squeeze() + self.running_var = state_dict["running_var"].squeeze() + self.scale = state_dict["weight"].squeeze() + self.bias = state_dict["bias"].squeeze() + elif isinstance(self.expr, CallFunction): + assert False, "not inplement function batchnorm" + self.op = BatchNormalizationOpr() + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + for i in self.expr.args[1:]: + t = self.resolver.get_ir_tensor(i, user_opr=self.op) + self.op.add_inp_tensors(t) + self.add_const_inputs() + elif isinstance(self.expr, CallFunction): + assert False, "not inplement function batchnorm" + + self.add_opr_out_tensors() + + def add_opr_out_tensors(self): + for o in self.expr.outputs: + out_tensor = self.resolver.get_ir_tensor(o, self.op) + self.op.add_out_tensors(out_tensor) + + def add_const_inputs(self): + if self.scale is not None: + scale_tensor = self.resolver.get_ir_tensor( + self.scale, + owner_opr=None, + name=self.expr.inputs[0]._name + "_scale", + user_opr=self.op, + ) + self.op.add_inp_tensors(scale_tensor) + if self.bias is not None: + bias_tensor = self.resolver.get_ir_tensor( + self.bias, + owner_opr=None, + name=self.expr.inputs[0]._name + "_bias", + user_opr=self.op, + ) + self.op.add_inp_tensors(bias_tensor) + + if self.running_mean is not None: + mean_tensor = self.resolver.get_ir_tensor( + self.running_mean, + owner_opr=None, + name=self.expr.inputs[0]._name + "_runing_mean", + user_opr=self.op, + ) + self.op.add_inp_tensors(mean_tensor) + + if self.running_var is not None: + var_tensor = self.resolver.get_ir_tensor( + self.running_var, + owner_opr=None, + name=self.expr.inputs[0]._name + "_runing_var", + user_opr=self.op, + ) + self.op.add_inp_tensors(var_tensor) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/broadcast.py b/mgeconvert/frontend/tm_to_ir/op_generators/broadcast.py new file mode 100644 index 0000000..4e7b69b --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/broadcast.py @@ -0,0 +1,39 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import numpy as np +from megengine.traced_module.expr import CallFunction + +from ....converter_ir.ir_op import BroadcastOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.broadcast_to) +class GenBroadcastOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + self.op = BroadcastOpr() + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + shape = self.expr.args[1] + if isinstance(shape, (list, tuple)): + shape = np.array(shape) + shape_tensor = self.resolver.get_ir_tensor(shape, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.op.add_inp_tensors(shape_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/concat.py b/mgeconvert/frontend/tm_to_ir/op_generators/concat.py new file mode 100644 index 0000000..0baed59 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/concat.py @@ -0,0 +1,80 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from collections import Iterable + +import megengine.functional as F +import megengine.module as M +import megengine.module.qat as QAT +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import ConcatOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.concat, M.Concat) +class GenConcatOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + if isinstance(self.expr, CallMethod): + self.axis = ( + self.expr.kwargs["axis"] + if "axis" in self.expr.kwargs + else self.expr.args[2] + ) + elif isinstance(self.expr, CallFunction): + self.axis = ( + self.expr.kwargs["axis"] + if "axis" in self.expr.kwargs + else self.expr.args[1] + ) + self.op = ConcatOpr(self.axis) + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + inp_data = self.expr.args[1] + elif isinstance(self.expr, CallFunction): + inp_data = self.expr.args[0] + assert isinstance(inp_data, Iterable), "Concat inputs must be Iterable." + for i in inp_data: + t = self.resolver.get_ir_tensor(i, user_opr=self.op) + self.op.add_inp_tensors(t) + self.add_opr_out_tensors() + + +@_register_op(QAT.Concat) +class GenQConcatOpr(GenConcatOpr): + def __init__(self, expr, irgraph) -> None: + if isinstance(expr, CallMethod): + self.module = expr.inputs[0].expr.value + if hasattr(self.module.act_fake_quant, "get_qparams"): + self.act_qparams = self.module.act_fake_quant.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + elif hasattr(self.module.act_observer, "get_qparams"): + self.act_qparams = self.module.act_observer.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + else: + logger.error("Observer and FakeQuantize do not have get_qparams().") + super().__init__(expr, irgraph=irgraph) + + def add_opr_out_tensors(self): + for o in self.expr.outputs: + t = self.resolver.get_ir_tensor(o, owner_opr=self.op) + t.set_qparams( + *self.resolver.resolve_qparams( + self.act_qparams.scale, self.act_qparams.zero_point + ) + ) + t.q_dtype = self.act_dtype + self.op.add_out_tensors(t) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/constant.py b/mgeconvert/frontend/tm_to_ir/op_generators/constant.py new file mode 100644 index 0000000..a705262 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/constant.py @@ -0,0 +1,22 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ....converter_ir.ir_op import ConstantOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op("Constant") +class GenConstantOpr(OpGenBase): + name = "Constant" + + def __init__(self, expr=None, net=None) -> None: + super().__init__(expr, net) + self.op = ConstantOpr() + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/conv2d.py b/mgeconvert/frontend/tm_to_ir/op_generators/conv2d.py new file mode 100644 index 0000000..35f5f70 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/conv2d.py @@ -0,0 +1,160 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from abc import ABC + +import megengine.functional as F +import megengine.module as M +import megengine.module.qat as QAT +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import Conv2dOpr +from ....converter_ir.ir_tensor import AxisOrder +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +class GenConvBase(OpGenBase, ABC): + def __init__(self, expr, irgraph, op_cls): + super().__init__(expr, irgraph) + if isinstance(expr, CallMethod): + conv_module = expr.inputs[0].expr.value + self.weight = conv_module.weight + self.bias = conv_module.bias + self.stride = conv_module.stride + self.padding = conv_module.padding + self.dilation = conv_module.dilation + self.groups = conv_module.groups + elif isinstance(expr, CallFunction): + self.weight = None + self.stride = self.expr.args[3] + self.padding = self.expr.args[4] + self.dilation = self.expr.args[5] + self.groups = self.expr.args[6] + if len(expr.args) > 7: + assert self.expr.args[7] == "cross_correlation" + if len(expr.args) > 8: + assert self.expr.args[8] == "default" + self.op = op_cls(self.stride, self.padding, self.dilation, self.groups,) + + def add_opr_vars(self, weight_format): + self.add_weight_bias_tensors(weight_format) + self.add_opr_out_tensors() + + def add_weight_bias_tensors(self, weight_format): + if isinstance(self.expr, CallMethod): + for i in self.expr.args[1:]: + t = self.resolver.get_ir_tensor(i, user_opr=self.op) + self.op.add_inp_tensors(t) + self.add_const_inputs(weight_format) + elif isinstance(self.expr, CallFunction): + inp_tensor = self.resolver.get_ir_tensor( + self.expr.args[0], user_opr=self.op + ) + self.op.add_inp_tensors(inp_tensor) + weight_tensor = self.resolver.get_ir_tensor( + self.expr.args[1], user_opr=self.op, axis_order=weight_format, + ) + weight_tensor.axis_order = weight_format + self.op.add_inp_tensors(weight_tensor) + if self.expr.args[2]: + bias = self.expr.args[2] + bias.shape = bias.shape[1] + bias_tensor = self.resolver.get_ir_tensor( + bias, name=self.expr.args[0]._name + "_bias", user_opr=self.op + ) + self.op.add_inp_tensors(bias_tensor) + + def add_const_inputs(self, weight_format): + if self.weight is not None: + weight_tensor = self.resolver.get_ir_tensor( + self.weight, + name=self.expr.inputs[0]._name + "_weight", + user_opr=self.op, + axis_order=weight_format, + ) + weight_tensor.axis_order = weight_format + self.op.add_inp_tensors(weight_tensor) + if self.bias is not None: + bias_tensor = self.resolver.get_ir_tensor( + self.bias, name=self.expr.inputs[0]._name + "_bias", user_opr=self.op, + ) + self.op.add_inp_tensors(bias_tensor) + + +@_register_op(M.Conv2d, F.conv2d) +class GenConv2dOpr(GenConvBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph, Conv2dOpr) + self.add_opr_vars(AxisOrder.OIHW) + + +class GenQConvBase(GenConvBase): + def __init__(self, expr, irgraph, op_cls): + conv_module = expr.inputs[0].expr.value + if hasattr(conv_module.weight_fake_quant, "get_qparams"): + self.weight_qparams = conv_module.weight_fake_quant.get_qparams() + self.weight_dtype = self.weight_qparams.dtype_meta.np_dtype_str + self.act_qparams = conv_module.act_fake_quant.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + elif hasattr(conv_module.weight_observer, "get_qparams"): + self.weight_qparams = conv_module.weight_observer.get_qparams() + self.weight_dtype = self.weight_qparams.dtype_meta.np_dtype_str + self.act_qparams = conv_module.act_observer.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + else: + logger.error("Observer and FakeQuantize do not have get_qparams().") + super().__init__(expr, irgraph, op_cls) + + def add_opr_out_tensors(self): + for o in self.expr.outputs: + out_tensor = self.resolver.get_ir_tensor(o, self.op) + out_tensor.set_qparams( + *self.resolver.resolve_qparams( + self.act_qparams.scale, self.act_qparams.zero_point + ) + ) + out_tensor.q_dtype = self.act_dtype + self.op.add_out_tensors(out_tensor) + + def add_const_inputs(self, weight_format): + if self.weight is not None: + weight_tensor = self.resolver.get_ir_tensor( + self.weight, + user_opr=self.op, + name=self.expr.inputs[0]._name + "_weight", + axis_order=weight_format, + ) + weight_tensor.set_qparams( + *self.resolver.resolve_qparams( + self.weight_qparams.scale, self.weight_qparams.zero_point + ) + ) + weight_tensor.q_dtype = self.weight_dtype + self.op.add_inp_tensors(weight_tensor) + + if self.bias is not None: + bias_tensor = self.resolver.get_ir_tensor( + self.bias, user_opr=self.op, name=self.expr.inputs[0]._name + "_bias" + ) + bias_tensor.set_qparams( + self.op.inp_tensors[0].scale * weight_tensor.scale, 0 + ) + bias_tensor.q_dtype = "int32" + self.op.add_inp_tensors(bias_tensor) + + +@_register_op(QAT.Conv2d) +class GenQConv2dOpr(GenQConvBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph, Conv2dOpr) + self.add_opr_vars(AxisOrder.OIHW) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/deconv.py b/mgeconvert/frontend/tm_to_ir/op_generators/deconv.py new file mode 100644 index 0000000..48e8d0f --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/deconv.py @@ -0,0 +1,35 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=no-member + +import megengine.functional as F +import megengine.module as M +import megengine.module.qat as QAT + +from ....converter_ir.ir_op import Deconv2dOpr +from ....converter_ir.ir_tensor import AxisOrder +from ..tm_utils import get_logger +from .base import _register_op +from .conv2d import GenConvBase, GenQConvBase + +logger = get_logger(__name__) + + +@_register_op(M.ConvTranspose2d, F.conv_transpose2d) +class GenDeconv2dOpr(GenConvBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph, Deconv2dOpr) + self.add_opr_vars(AxisOrder.IOHW) + + +@_register_op(QAT.ConvTranspose2d) +class GenQDeconv2dOpr(GenQConvBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph, Deconv2dOpr) + self.add_opr_vars(AxisOrder.IOHW) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/dropout.py b/mgeconvert/frontend/tm_to_ir/op_generators/dropout.py new file mode 100644 index 0000000..142c744 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/dropout.py @@ -0,0 +1,44 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import megengine.module as M +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import DropoutOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(M.Dropout, F.dropout) +class GenDropoutOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + if isinstance(expr, CallMethod): + if hasattr(expr.kwargs, "drop_prob"): + self.drop_prob = expr.kwargs["drop_prob"] + else: + self.drop_prob = 0 + self.training = False + + if isinstance(expr, CallFunction): + assert False, "functional.dropout is not implement" + + self.op = DropoutOpr(self.drop_prob, self.training) + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + for i in self.expr.args[1:]: + t = self.resolver.get_ir_tensor(i, user_opr=self.op) + self.op.add_inp_tensors(t) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/elemwise.py b/mgeconvert/frontend/tm_to_ir/op_generators/elemwise.py new file mode 100644 index 0000000..19f3baf --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/elemwise.py @@ -0,0 +1,288 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +from abc import ABC + +import megengine.functional as F +import megengine.module as M +import megengine.module.qat as QAT +import numpy as np +from megengine.traced_module.expr import CallFunction, CallMethod, Constant +from megengine.traced_module.node import ModuleNode, TensorNode + +from ....converter_ir.ir_op import ( + AbsOpr, + AddOpr, + CeilOpr, + ExpOpr, + FloorDivOpr, + FloorOpr, + HardSigmoidOpr, + HardSwishOpr, + IdentityOpr, + LeakyReluOpr, + LogOpr, + MaxOpr, + MinOpr, + MulOpr, + PowOpr, + Relu6Opr, + ReluOpr, + SigmoidOpr, + SiLUOpr, + SubOpr, + TanHOpr, + TrueDivOpr, +) +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.leaky_relu) +class GenLeakyReluOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + if len(expr.const_val) == 0: + self.op = LeakyReluOpr() + else: + self.negative_slope = expr.args[1] + self.op = LeakyReluOpr(self.negative_slope) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() + + +class GenElemwiseOpr(OpGenBase, ABC): + def __init__(self, expr, irgraph, op_cls): + super().__init__(expr, irgraph) + self.op = op_cls() + self.add_opr_vars() + + def add_opr_vars(self): + for inp in self.expr.args: + if isinstance(inp, ModuleNode): + continue + assert isinstance( + inp, (TensorNode, int, float, np.ndarray) + ), "expr inputs type not support {}".format(type(inp)) + tm_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + if isinstance(inp, (int, float, np.ndarray)): + target_dtype = self.op.inp_tensors[0].dtype + tm_tensor.set_dtype(target_dtype) + self.op.add_inp_tensors(tm_tensor) + + for oup in self.expr.outputs: + assert isinstance( + oup, (TensorNode) + ), "expr outputs type not support {}".format(type(oup)) + tm_tensor = self.resolver.get_ir_tensor(oup, owner_opr=self.op) + self.op.add_out_tensors(tm_tensor) + if ( + hasattr(self.op.inp_tensors[0], "scale") + and self.op.inp_tensors[0].scale is not None + ): + for o in self.op.out_tensors: + o.scale = self.op.inp_tensors[0].scale + o.zero_point = self.op.inp_tensors[0].zero_point + o.dtype = self.op.inp_tensors[0].dtype + # set dtype for const value + + +@_register_op("__add__", "__iadd__") +class GenAddOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, AddOpr) + + +@_register_op("__sub__") +class GenSubOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, SubOpr) + + +@_register_op("__mul__") +class GenMulOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, MulOpr) + + +@_register_op("__truediv__") +class GenTrueDivOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, TrueDivOpr) + + +@_register_op("__floordiv__") +class GenFloorDivOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, FloorDivOpr) + + +@_register_op("__pow__") +class GenPowOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, PowOpr) + + +@_register_op(F.maximum) +class GenMaxOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, MaxOpr) + + +@_register_op(F.minimum) +class GenMinOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, MinOpr) + + +@_register_op(F.exp) +class GenExpOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, ExpOpr) + + +@_register_op(F.floor) +class GenFloorOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, FloorOpr) + + +@_register_op(F.ceil) +class GenCeilOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, CeilOpr) + + +@_register_op(F.abs) +class GenAbsOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, AbsOpr) + + +@_register_op(F.relu, M.activation.ReLU) +class GenReluOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, ReluOpr) + + +@_register_op(F.tanh) +class GenTanHOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, TanHOpr) + + +@_register_op(F.sigmoid, M.Sigmoid) +class GenSigmoidOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, SigmoidOpr) + + +@_register_op(F.hsigmoid) +class GenHardSigmoidOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, HardSigmoidOpr) + self.add_const_vars() + + def add_const_vars(self): + const_tensor3 = self.resolver.get_ir_tensor(3, user_opr=self.op) + const_tensor6 = self.resolver.get_ir_tensor(6, user_opr=self.op) + self.op.add_inp_tensors(const_tensor3) + self.op.add_inp_tensors(const_tensor6) + + +@_register_op(F.log) +class GenLogOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, LogOpr) + + +@_register_op(F.relu6) +class GenRelu6Opr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, Relu6Opr) + + +@_register_op(F.silu, M.SiLU) +class GenSiLUOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, SiLUOpr) + + +@_register_op(F.hswish) +class GenHardSwishOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, HardSwishOpr) + self.add_const_vars() + + def add_const_vars(self): + const_tensor3 = self.resolver.get_ir_tensor(3, user_opr=self.op) + const_tensor6 = self.resolver.get_ir_tensor(6, user_opr=self.op) + self.op.add_inp_tensors(const_tensor3) + self.op.add_inp_tensors(const_tensor6) + + +@_register_op(M.Identity) +class GenIdentityOpr(GenElemwiseOpr): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph, IdentityOpr) + + +method_opr_map = { + "add": GenAddOpr, + "sigmoid": GenSigmoidOpr, + "mul": GenMulOpr, + "abs": GenAbsOpr, + "ceil": GenCeilOpr, + "exp": GenExpOpr, + "floor": GenFloorOpr, + "log": GenLogOpr, + "max": GenMaxOpr, + "min": GenMinOpr, + "pow": GenPowOpr, + "relu": GenReluOpr, + "sub": GenSubOpr, + "tanh": GenTanHOpr, + "true_div": GenTrueDivOpr, + "floor_div": GenFloorDivOpr, + "relu6": GenRelu6Opr, + "identity": GenIdentityOpr, + "hswish": GenHardSwishOpr, + "hsigmoid": GenHardSigmoidOpr, +} + + +@_register_op(QAT.Elemwise, M.Elemwise) +def get_elemwise_op(expr, net): + assert isinstance(expr, CallMethod) + assert isinstance(expr.inputs[0].expr, Constant) + module = expr.inputs[0].expr.value + method = module.method + op_gen = method_opr_map[method](expr, net) + + if isinstance(module, QAT.QATModule): + if hasattr(module.act_fake_quant, "get_qparams"): + qparams = module.act_fake_quant.get_qparams() + else: + qparams = module.act_observer.get_qparams() + for o in op_gen.get_opr().out_tensors: + o.scale = float(qparams.scale) if method != "sigmoid" else 1 / 256.0 + o.zero_point = ( + int(qparams.zero_point) if qparams.zero_point is not None else None + ) + o.q_dtype = qparams.dtype_meta.np_dtype_str + return op_gen diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/flatten.py b/mgeconvert/frontend/tm_to_ir/op_generators/flatten.py new file mode 100644 index 0000000..020f46a --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/flatten.py @@ -0,0 +1,39 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction + +from ....converter_ir.ir_op import FlattenOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.flatten) +class GenFlattenOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + self.start_axis = 0 + self.end_axis = -1 + if len(expr.args) >= 2: + self.start_axis = int(expr.args[1]) + if len(expr.args) >= 3: + self.end_axis = int(expr.args[2]) + self.op = FlattenOpr(self.start_axis, self.end_axis) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/getvarshape.py b/mgeconvert/frontend/tm_to_ir/op_generators/getvarshape.py new file mode 100644 index 0000000..f8c7db7 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/getvarshape.py @@ -0,0 +1,26 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from ....converter_ir.ir_op import GetVarShapeOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op("GetVarShape") +class GenGetVarShapeOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + self.op = GetVarShapeOpr() + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/matmul.py b/mgeconvert/frontend/tm_to_ir/op_generators/matmul.py new file mode 100644 index 0000000..18f479d --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/matmul.py @@ -0,0 +1,134 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import megengine.module as M +import megengine.module.qat as QAT +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import LinearOpr, MatMulOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.matmul) +class GenMatMulOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + if isinstance(expr, CallFunction): + self.transpose_a = expr.args[2] + self.transpose_b = expr.args[3] + self.compute_mode = expr.args[4] + self.format = expr.args[5] + self.op = MatMulOpr( + self.transpose_a, self.transpose_b, self.compute_mode, self.format + ) + self.add_opr_vars() + + def add_opr_vars(self): + for inp in self.expr.inputs: + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() + + +@_register_op(F.linear, M.Linear) +class GenLinearOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + if isinstance(expr, CallMethod): + m = expr.inputs[0].expr.value + self.weight = m.weight + self.has_bias = bool(m.bias is not None) + elif isinstance(expr, CallFunction): + self.has_bias = bool(len(expr.inputs) == 3) + self.op = LinearOpr(self.has_bias) + + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + for inp in self.expr.inputs[1:]: + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_const_inputs() + else: + for inp in self.expr.inputs: + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() + + def add_const_inputs(self): + weight_tensor = self.resolver.get_ir_tensor( + self.weight, user_opr=self.op, name=self.expr.inputs[0]._name + "_weight", + ) + self.op.add_inp_tensors(weight_tensor) + if self.has_bias: + bias_tensor = self.resolver.get_ir_tensor( + self.expr.inputs[0].expr.value.bias, + user_opr=self.op, + name=self.expr.inputs[0]._name + "_bias", + ) + self.op.add_inp_tensors(bias_tensor) + + +@_register_op(QAT.Linear) +class GenQLinearOpr(GenLinearOpr): + def __init__(self, expr, irgraph): + self.module = expr.inputs[0].expr.value + if hasattr(self.module.weight_fake_quant, "get_qparams"): + self.weight_qparams = self.module.weight_fake_quant.get_qparams() + self.weight_dtype = self.weight_qparams.dtype_meta.np_dtype_str + self.act_qparams = self.module.act_fake_quant.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + elif hasattr(self.module.weight_observer, "get_qparams"): + self.weight_qparams = self.module.weight_observer.get_qparams() + self.weight_dtype = self.weight_qparams.dtype_meta.np_dtype_str + self.act_qparams = self.module.act_observer.get_qparams() + self.act_dtype = self.act_qparams.dtype_meta.np_dtype_str + else: + logger.error("Observer and FakeQuantize do not have get_qparams().") + super().__init__(expr, irgraph) + + def add_const_inputs(self): + weight_tensor = self.resolver.get_ir_tensor( + self.weight, user_opr=self.op, name=self.expr.inputs[0]._name + "_weight", + ) + weight_tensor.set_qparams( + *self.resolver.resolve_qparams( + self.weight_qparams.scale, self.weight_qparams.zero_point + ) + ) + weight_tensor.q_dtype = self.weight_dtype + self.op.add_inp_tensors(weight_tensor) + if self.has_bias: + bias_tensor = self.resolver.get_ir_tensor( + self.expr.inputs[0].expr.value.bias, + user_opr=self.op, + name=self.expr.inputs[0]._name + "_bias", + ) + bias_tensor.set_qparams( + self.op.inp_tensors[0].scale * weight_tensor.scale, 0 + ) + bias_tensor.q_dtype = "int32" + self.op.add_inp_tensors(bias_tensor) + + def add_opr_out_tensors(self): + for o in self.expr.outputs: + out_tensor = self.resolver.get_ir_tensor(o, owner_opr=self.op) + out_tensor.set_qparams( + *self.resolver.resolve_qparams( + self.act_qparams.scale, self.act_qparams.zero_point + ) + ) + out_tensor.q_dtype = self.act_dtype + self.op.add_out_tensors(out_tensor) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/pooling.py b/mgeconvert/frontend/tm_to_ir/op_generators/pooling.py new file mode 100644 index 0000000..0016404 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/pooling.py @@ -0,0 +1,89 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import megengine.module as M +from megengine.traced_module.expr import CallFunction, CallMethod +from megengine.utils.tuple_function import _pair, _pair_nonzero + +from ....converter_ir.ir_op import AdaptiveAvgPool2dOpr, AvgPool2dOpr, MaxPool2dOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + +mode_map = { + M.MaxPool2d: MaxPool2dOpr, + F.max_pool2d: MaxPool2dOpr, + M.AvgPool2d: AvgPool2dOpr, + F.avg_pool2d: AvgPool2dOpr, +} + + +@_register_op( + M.MaxPool2d, + F.max_pool2d, + M.AvgPool2d, + F.avg_pool2d, + M.adaptive_pooling.AdaptiveAvgPool2d, + F.adaptive_avg_pool2d, +) +class GenPool2dOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + + if isinstance(expr, CallMethod): + m = expr.inputs[0].expr.value + self.kernel_size = _pair_nonzero(m.kernel_size) + self.stride = _pair_nonzero(m.stride) + self.padding = _pair(m.padding) + op_cls = mode_map[type(m)] + elif isinstance(expr, CallFunction): + self.kernel_size = self.expr.args[1] + self.stride = self.kernel_size + self.padding = 0 + if len(self.expr.args) > 2: + self.stride = self.expr.args[2] + if len(self.expr.args) > 3: + self.padding = self.expr.args[3] + op_cls = mode_map[expr.func] + self.op = op_cls(self.kernel_size, self.stride, self.padding) + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + inp = self.expr.args[1] + elif isinstance(self.expr, CallFunction): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() + + +@_register_op(M.adaptive_pooling.AdaptiveAvgPool2d, F.adaptive_avg_pool2d) +class GenAdaptiveAvgPool2dOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + + if isinstance(expr, CallMethod): + m = expr.inputs[0].expr.value + self.op = AdaptiveAvgPool2dOpr(m.oshp) + elif isinstance(expr, CallFunction): + self.op = AdaptiveAvgPool2dOpr(expr.func.oshp) + self.add_opr_vars() + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + inp = self.expr.args[1] + elif isinstance(self.expr, CallFunction): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/reduce.py b/mgeconvert/frontend/tm_to_ir/op_generators/reduce.py new file mode 100644 index 0000000..d183832 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/reduce.py @@ -0,0 +1,42 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction + +from ....converter_ir.ir_op import ReduceOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + +func_mode_map = { + F.max: "MAX", + F.min: "MIN", + F.mean: "MEAN", + F.sum: "SUM", +} + + +@_register_op(F.max, F.min, F.mean, F.sum) +class GenReduceOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + self.axis = expr.kwargs["axis"] + self.mode = func_mode_map[expr.func] + self.op = ReduceOpr(self.axis, self.mode, False) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/reshape.py b/mgeconvert/frontend/tm_to_ir/op_generators/reshape.py new file mode 100644 index 0000000..a424ab1 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/reshape.py @@ -0,0 +1,61 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import RepeatOpr, ReshapeOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.reshape, "reshape") +class GenReshapeOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + if isinstance(self.expr, CallFunction): + self.out_shape = expr.args[1] + self.op = ReshapeOpr(self.out_shape) + elif isinstance(self.expr, CallMethod): + self.out_shape = self.expr.outputs[0].shape + self.op = ReshapeOpr(self.out_shape) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + if isinstance(self.expr, CallMethod): + if len(self.expr.args) > 1: + for inp in self.expr.args[1:]: + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() + + +@_register_op(F.repeat) +class GenRepeatOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + self.repeats = self.expr.args[1] + self.axis = None + if "axis" in self.expr.kwargs.keys(): + self.axis = self.expr.kwargs["axis"] + self.op = RepeatOpr(self.repeats, self.axis) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/resize.py b/mgeconvert/frontend/tm_to_ir/op_generators/resize.py new file mode 100644 index 0000000..c253f00 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/resize.py @@ -0,0 +1,45 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction + +from ....converter_ir.ir_op import ResizeOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.vision.interpolate) +class GenResizeOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + assert isinstance(expr, CallFunction) + self.size = expr.kwargs["size"] + self.scale_factor = ( + expr.kwargs["scale_factor"] + if "scale_factor" in expr.kwargs.keys() + else None + ) + self.mode = expr.kwargs["mode"] if "mode" in expr.kwargs.keys() else None + self.align_corners = ( + expr.kwargs["align_corners"] + if "align_corners" in expr.kwargs.keys() + else False + ) + self.op = ResizeOpr(self.size, self.scale_factor, self.mode, self.align_corners) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/softmax.py b/mgeconvert/frontend/tm_to_ir/op_generators/softmax.py new file mode 100644 index 0000000..ec94348 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/softmax.py @@ -0,0 +1,82 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +import megengine.module as M +import numpy as np +from megengine.traced_module.expr import CallFunction, CallMethod +from mgeconvert.converter_ir.ir_tensor import IRTensor + +from ....converter_ir.ir_op import SoftmaxOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +def get_softmax_axis(ndim: int) -> int: + if ndim in (0, 1, 3): + return 0 + return 1 + + +@_register_op(M.Softmax, F.softmax) +class GenSoftmaxOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + if isinstance(self.expr, CallMethod): + module = expr.inputs[0].expr.value + self.axis = module.axis + elif isinstance(self.expr, CallFunction): + if len(expr.const_val) == 1: + self.axis = expr.const_val + else: + self.axis = None + + self.op = SoftmaxOpr(self.axis) + self.add_opr_vars() + assert self.op.inp_tensors[0].ndim in [ + 1, + 2, + 4, + ], "Softmax do not support {} dim".format(self.op.inp_tensors[0].ndim) + if self.op.axis is None: + self.op.axis = get_softmax_axis(self.op.inp_tensors[0].ndim) + + if self.op.axis is None: + self.op.axis = get_softmax_axis(self.op.inp_tensors[0].ndim) + + def add_opr_vars(self): + if isinstance(self.expr, CallMethod): + inp_tensor = self.resolver.get_ir_tensor( + self.expr.args[1], user_opr=self.op + ) + self.op.add_inp_tensors(inp_tensor) + + elif isinstance(self.expr, CallFunction): + inp_tensor = self.resolver.get_ir_tensor( + self.expr.args[0], user_opr=self.op + ) + self.op.add_inp_tensors(inp_tensor) + + # self.add_axis() + self.add_opr_out_tensors() + + def add_axis(self): + if self.axis is None: + self.axis = get_softmax_axis(self.op.inp_tensors[0].ndim) + axis_tensor = IRTensor( + self.op.inp_tensors[0].name + "_softmax_axis", + (1,), + np.int32, + np_data=np.array(self.axis), + axis=None, + ) + self.op.add_inp_tensors(axis_tensor) diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/squeeze.py b/mgeconvert/frontend/tm_to_ir/op_generators/squeeze.py new file mode 100644 index 0000000..1cdbe12 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/squeeze.py @@ -0,0 +1,38 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from typing import Sequence + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction + +from ....converter_ir.ir_op import SqueezeOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op(F.squeeze) +class GenSqueezeOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallFunction) + self.squeeze_dims = expr.args[1] + if not isinstance(self.squeeze_dims, Sequence): + self.squeeze_dims = (self.squeeze_dims,) + self.op = SqueezeOpr(self.squeeze_dims) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/subtensor.py b/mgeconvert/frontend/tm_to_ir/op_generators/subtensor.py new file mode 100644 index 0000000..a9bb668 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/subtensor.py @@ -0,0 +1,77 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from typing import List # pylint: disable=unused-import +from typing import Sequence + +import numpy as np +from megengine.traced_module.expr import CallMethod + +from ....converter_ir.ir_op import GetSubTensorOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op("__getitem__") +class GenGetSubtensorOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallMethod) + self.axis = [] # type: List[int] + self.begin_params = [] # type: List[int] + self.end_params = [] # type: List[int] + self.step_params = [] # type: List[int] + self.squeeze_axis = [] # type: List[int] + if isinstance(expr.args[1], Sequence): + for i in range(len(expr.args[1])): + slice = expr.args[1][i] + if isinstance(slice, int): + start = slice + stop = slice + 1 + step = 1 + self.squeeze_axis.append(i) + elif slice.start or slice.stop: + start = slice.start if slice.start is not None else 0 + stop = ( + slice.stop if slice.stop is not None else np.iinfo(np.int32).max + ) + step = slice.step if slice.step is not None else 1 + else: + continue + self.begin_params.append(start) + self.end_params.append(stop) + self.step_params.append(step) + self.axis.append(i) + elif isinstance(expr.args[1], int): + start = expr.args[1] + stop = expr.args[1] + 1 + step = 1 + self.squeeze_axis.append(0) + self.begin_params.append(start) + self.end_params.append(stop) + self.step_params.append(step) + self.axis.append(0) + + self.op = GetSubTensorOpr( + self.axis, + self.begin_params, + self.end_params, + self.step_params, + self.squeeze_axis, + ) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/transpose.py b/mgeconvert/frontend/tm_to_ir/op_generators/transpose.py new file mode 100644 index 0000000..5bc63a1 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/transpose.py @@ -0,0 +1,33 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +from megengine.traced_module.expr import CallFunction, CallMethod + +from ....converter_ir.ir_op import TransposeOpr +from .base import OpGenBase, _register_op + + +@_register_op(F.transpose, "transpose") +class GenTransposeOpr(OpGenBase): + def __init__(self, expr, irgraph) -> None: + super().__init__(expr, irgraph) + if isinstance(self.expr, CallFunction): + self.pattern = expr.args[1] + elif isinstance(self.expr, CallMethod): + self.pattern = tuple(expr.args[1:]) + self.op = TransposeOpr(self.pattern) + self.add_opr_vars() + + def add_opr_vars(self): + inp_tensor = self.resolver.get_ir_tensor(self.expr.args[0], user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/op_generators/typecvt.py b/mgeconvert/frontend/tm_to_ir/op_generators/typecvt.py new file mode 100644 index 0000000..b2f394b --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/op_generators/typecvt.py @@ -0,0 +1,33 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from megengine.traced_module.expr import CallMethod + +from ....converter_ir.ir_op import TypeCvtOpr +from ..tm_utils import get_logger +from .base import OpGenBase, _register_op + +logger = get_logger(__name__) + + +@_register_op("astype") +class GenTypeCvtOpr(OpGenBase): + def __init__(self, expr, irgraph): + super().__init__(expr, irgraph) + assert isinstance(self.expr, CallMethod) + out_dtype = expr.args[1] + self.op = TypeCvtOpr(out_dtype) + self.add_opr_vars() + + def add_opr_vars(self): + inp = self.expr.args[0] + inp_tensor = self.resolver.get_ir_tensor(inp, user_opr=self.op) + self.op.add_inp_tensors(inp_tensor) + self.add_opr_out_tensors() diff --git a/mgeconvert/frontend/tm_to_ir/pattern_utils.py b/mgeconvert/frontend/tm_to_ir/pattern_utils.py new file mode 100644 index 0000000..65e2cf3 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/pattern_utils.py @@ -0,0 +1,161 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +from collections import OrderedDict +from functools import partial +from typing import Union + +from megengine.logger import get_logger +from megengine.module import Module +from megengine.traced_module.expr import ( + Apply, + CallFunction, + CallMethod, + Constant, + Expr, + Input, +) +from megengine.traced_module.node import ModuleNode, Node +from megengine.traced_module.pytree import SUPPORTED_TYPE, LeafDef, tree_flatten + +logger = get_logger(__name__) + +DEFAULT_FUSION_PATTERNS = OrderedDict() # type: OrderedDict + + +def flatten(treedef, values): + if isinstance(treedef, LeafDef): + return [ + values, + ] + if type(values) != treedef.type: # pylint:disable=unidiomatic-typecheck + return None + rst = [] + children_values, aux_data = SUPPORTED_TYPE[treedef.type].flatten(values) + if treedef.aux_data != aux_data: + return None + for ch_def, ch_val in zip(treedef.children_defs, children_values): + v_list = flatten(ch_def, ch_val) + if v_list is None: + return None + rst.extend(v_list) + if treedef.num_leaves != len(rst): + return None + return rst + + +class MatchAnyNode: + pass + + +class InputNode: + pass + + +def register_pattern(pattern, default_dict: OrderedDict): + def insert(func): + default_dict[pattern] = func + return func + + return insert + + +register_fusion_pattern = partial( + register_pattern, default_dict=DEFAULT_FUSION_PATTERNS +) + + +def check_match(expr: Expr, target): # pylint:disable=too-many-return-statements + if isinstance(target, type) and issubclass(target, MatchAnyNode): + return True + + if target is InputNode and isinstance(expr, Input): + return True + + if isinstance(expr, Apply): + opdef = expr.opdef + return isinstance(opdef, target) + + if isinstance(target, type) and issubclass(target, Module): + if not isinstance(expr, (CallMethod, Constant)): + return False + if isinstance(expr, CallMethod): + obj_node = expr.inputs[0] + if not isinstance(obj_node, ModuleNode): + return False + obj = obj_node.owner + else: + obj = expr.value + if type(obj) != target: # pylint:disable=unidiomatic-typecheck + return False + elif callable(target): + if not isinstance(expr, CallFunction): + return False + if expr.func != target: + return False + elif isinstance(target, str): + if not isinstance(expr, CallMethod): + return False + if expr.method != target: + return False + else: + if expr != target: + return False + + return True + + +def is_match( + expr: Union[Expr, Node], pattern: tuple, max_uses=100 +): # pylint:disable=too-many-return-statements + + if isinstance(pattern, tuple): + self_match, *arg_matches = pattern + else: + self_match = pattern + arg_matches = [] + + if isinstance(expr, Node): + expr = expr.expr + + if isinstance(expr, Expr): + if max_uses == 1: + if len(expr.outputs) > 1: + logger.warning("is_match only support 1 output expr") + return False + for n in expr.outputs: + if len(n.users) > max_uses: + return False + + if not check_match(expr, self_match): + return False + + if not arg_matches: + return True + + expr_args = expr.inputs + + if isinstance(expr, CallMethod) and isinstance(expr_args[0], ModuleNode): + expr_args = expr_args[1:] + if len(expr_args) != len(arg_matches): + return False + + for inp_node, arg_match in zip(expr_args, arg_matches): + inp_node, inp_def = tree_flatten(inp_node) + inp_match = flatten(inp_def, arg_match) + if inp_match is None: + return False + for node, arg in zip(inp_node, inp_match): + if isinstance(node, ModuleNode): + continue + if not is_match(node, arg, max_uses=1): + return False + + return True diff --git a/mgeconvert/frontend/tm_to_ir/qat_pattern.py b/mgeconvert/frontend/tm_to_ir/qat_pattern.py new file mode 100644 index 0000000..8cb8acc --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/qat_pattern.py @@ -0,0 +1,186 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine.functional as F +from megengine.core.ops import builtin +from megengine.module.qat import QATModule +from megengine.traced_module.expr import Apply, CallFunction, GetAttr + +from ...converter_ir.ir_op import ReluOpr +from .op_generators import GenConv2dOpr, GenDeconv2dOpr +from .pattern_utils import InputNode, MatchAnyNode, is_match, register_fusion_pattern +from .tm_tensor_resolver import TensorNodeResolver + + +def gen_qat_conv_opr(module, conv_function_expr, qat_expr, irgraph, is_deconv=False): + weight_fakequant = conv_function_expr.inputs[1].expr + bias = None + if len(conv_function_expr.inputs) == 3: + bias = conv_function_expr.inputs[2] + assert (isinstance(bias.expr, GetAttr) and bias.expr.name == "bias") or ( + isinstance(bias.expr, Apply) + and isinstance( + bias.expr.opdef, builtin.FakeQuant # pylint:disable=no-member + ) + ) + assert ( + isinstance(weight_fakequant.inputs[1].expr, GetAttr) + and weight_fakequant.inputs[1].expr.name == "weight" + ) + assert len(module.graph.inputs) == 2 + + act_qparams = module.act_fake_quant.get_qparams() + weight_qparams = module.weight_fake_quant.get_qparams() + + module.stride = conv_function_expr.args[3] + module.padding = conv_function_expr.args[4] + module.dilation = conv_function_expr.args[5] + module.groups = conv_function_expr.args[6] + assert conv_function_expr.args[7] == "cross_correlation" + assert conv_function_expr.args[8] == "default" + if bias is None: + module.bias = None + + op = ( + GenConv2dOpr(qat_expr, irgraph).get_opr() + if not is_deconv + else GenDeconv2dOpr(qat_expr, irgraph).get_opr() + ) + + op.inp_tensors[1].scale = float(weight_qparams.scale) + op.inp_tensors[1].zero_point = int(weight_qparams.zero_point) + op.inp_tensors[1].q_dtype = weight_qparams.dtype_meta.np_dtype_str + if len(op.inp_tensors) == 3: + op.inp_tensors[2].scale = op.inp_tensors[0].scale * op.inp_tensors[1].scale + op.inp_tensors[2].q_dtype = "int32" + op.inp_tensors[2].zero_point = 0 + + op.out_tensors[0].scale = act_qparams.scale.numpy()[0] + op.out_tensors[0].zero_point = act_qparams.zero_point.numpy()[0] + op.out_tensors[0].q_dtype = act_qparams.dtype_meta.np_dtype_str + return op + + +MATCH_RULE = {} + +pat_conv_bias_relu = ( + QATModule._apply_fakequant_with_observer, + MatchAnyNode, + ( + F.relu, + (F.conv2d, InputNode, QATModule._apply_fakequant_with_observer, MatchAnyNode), + ), + MatchAnyNode, +) + +pat_conv_bias = ( + QATModule._apply_fakequant_with_observer, + MatchAnyNode, + (F.conv2d, InputNode, QATModule._apply_fakequant_with_observer, MatchAnyNode), + MatchAnyNode, +) + +pat_conv_relu = ( + QATModule._apply_fakequant_with_observer, + MatchAnyNode, + (F.relu, (F.conv2d, InputNode, QATModule._apply_fakequant_with_observer),), + MatchAnyNode, +) + +pat_conv = ( + QATModule._apply_fakequant_with_observer, + MatchAnyNode, + (F.conv2d, InputNode, QATModule._apply_fakequant_with_observer), + MatchAnyNode, +) + +pat_deconv_relu = ( + QATModule._apply_fakequant_with_observer, + MatchAnyNode, + (F.relu, (F.conv_transpose2d, InputNode, QATModule._apply_fakequant_with_observer)), + MatchAnyNode, +) + + +@register_fusion_pattern(pat_conv_bias_relu) +def qat_conv_bias_relu(module, expr, call_expr, irgraph, _): + relu = expr.inputs[1].expr + op = gen_qat_conv_opr(module, relu.inputs[0].expr, call_expr, irgraph) + op.activation = "RELU" + return op + + +@register_fusion_pattern(pat_conv_bias) +def qat_conv_bias(module, expr, call_expr, irgraph, _): + conv = expr.inputs[1].expr + op = gen_qat_conv_opr(module, conv, call_expr, irgraph) + return op + + +@register_fusion_pattern(pat_conv_relu) +def qat_conv_relu(module, expr, call_expr, net, _): + relu = expr.inputs[1].expr + op = gen_qat_conv_opr(module, relu.inputs[0].expr, call_expr, net) + op.activation = "RELU" + return op + + +@register_fusion_pattern(pat_conv) +def qat_conv(module, expr, call_expr, net, _): + conv = expr.inputs[1].expr + op = gen_qat_conv_opr(module, conv, call_expr, net) + return op + + +@register_fusion_pattern(pat_deconv_relu) +def qat_deconv_relu_bias( + module, expr, call_expr, irgraph, resolver: TensorNodeResolver +): + relu = expr.inputs[1].expr + deconv = relu.inputs[0].expr + op = gen_qat_conv_opr(module, deconv, call_expr, irgraph, is_deconv=True) + op.activation = "RELU" + + relu_op = ReluOpr() + relu_op.inp_tensors = [] + relu_op.out_tensors = [] + relu_op.inp_tensors.append(op.out_tensors[0]) + relu_op.out_tensors.append(resolver.resolve(call_expr.outputs[0], relu_op)[0]) + relu_op.out_tensors[0].name += "_relu" + + relu_op.out_tensors[0].q_dtype = relu_op.inp_tensors[0].q_dtype + relu_op.out_tensors[0].scale = relu_op.inp_tensors[0].scale + relu_op.out_tensors[0].zero_point = relu_op.inp_tensors[0].zero_point + irgraph.all_tensors[ + irgraph._tensor_ids.index(call_expr.outputs[0]._id) + ] = relu_op.out_tensors[0] + + return op, relu_op + + +MATCH_RULE[QATModule._apply_fakequant_with_observer] = [ + pat_conv_bias_relu, + pat_conv_bias, + pat_deconv_relu, + pat_conv_relu, + pat_conv, +] + + +def find_match_pattern(graph): + rst = [] + for expr in graph._exprs: + if isinstance(expr, CallFunction): + if expr.func in MATCH_RULE: + pat = MATCH_RULE[expr.func] + for p in pat: + if is_match(expr, p): + rst.append((p, expr)) + return rst diff --git a/mgeconvert/frontend/tm_to_ir/tm_frontend.py b/mgeconvert/frontend/tm_to_ir/tm_frontend.py new file mode 100644 index 0000000..7193afa --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/tm_frontend.py @@ -0,0 +1,154 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint:disable=no-name-in-module,import-error +from typing import List, Sequence + +import megengine +from megengine.core._imperative_rt.core2 import Tensor as RawTensor +from megengine.module.qat import QuantStub +from megengine.traced_module import TracedModule +from megengine.traced_module.expr import ( + Apply, + CallFunction, + CallMethod, + Constant, + GetAttr, + Input, +) +from megengine.traced_module.module_tracer import BUILTIN_ARRAY_METHOD +from megengine.traced_module.node import ModuleNode, TensorNode + +from ...converter_ir.ir_graph import IRGraph +from .op_generators import EXPR2OP +from .pattern_utils import DEFAULT_FUSION_PATTERNS +from .qat_pattern import find_match_pattern +from .tm_tensor_resolver import TensorNodeResolver +from .tm_utils import get_logger + +logger = get_logger(__name__) + + +class TM_FrontEnd: + def __init__(self, traced_module): + if isinstance(traced_module, TracedModule): + self.module = traced_module.flatten() + elif isinstance(traced_module, str): + self.module = megengine.load(traced_module) + self.inputs: List[TensorNode] = self.module.graph.inputs[1:] + self.outputs: List[TensorNode] = self.module.graph.outputs + + self.irgraph = IRGraph() + self.tensor_resolver = TensorNodeResolver(self.irgraph) + + def resolve(self): + self.add_net_inputs() + self.get_all_oprs() + self.add_net_outputs() + return self.irgraph + + def add_net_inputs(self): + for node in self.inputs: + inp_tensor = self.tensor_resolver.get_ir_tensor(node, owner_opr=self) + if node.qparams is not None: + inp_tensor.set_qparams( + *self.tensor_resolver.resolve_qparams( + node.qparams.scale, node.qparams.zero_point + ) + ) + inp_tensor.q_dtype = node.qparams.dtype_meta.np_dtype_str + self.irgraph.add_net_inputs(inp_tensor) + + def get_all_oprs(self): + for expr in self.module.graph._exprs: + if isinstance(expr, Constant): + if isinstance(expr.value, RawTensor): + op_gen_cls = EXPR2OP.get("Constant") + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif isinstance(expr, GetAttr): + if isinstance(expr.outputs[0], TensorNode): + op_gen_cls = EXPR2OP.get("Constant") + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif isinstance(expr, CallMethod): + if expr.method in BUILTIN_ARRAY_METHOD: + # generate array_method op + op_gen_cls = EXPR2OP.get(expr.method, None) + assert op_gen_cls, "METHOD {} is not supported.".format(expr.method) + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif expr.method == "__new__": + # TODO + pass + elif expr.method == "__call__": + m = expr.inputs[0] + assert isinstance(m, ModuleNode) + assert isinstance(m.expr, Constant) + if isinstance(m.expr.value, TracedModule): + module = m.expr.value + assert module.is_qat + pats = find_match_pattern(module.graph) + pat, end_expr = pats[0] + fusion_op = DEFAULT_FUSION_PATTERNS.get(pat) + ops = fusion_op( + module, end_expr, expr, self.irgraph, self.tensor_resolver, + ) + ops = (ops,) if not isinstance(ops, Sequence) else ops + # if len(ops) >1: + # import pdb;pdb.set_trace() + for op in ops: + self.irgraph.all_oprs.append(op) + self.irgraph._opr_ids.append(id(op)) + elif isinstance(m.expr.value, QuantStub): + module = m.expr.value + inp_tensor = self.tensor_resolver.get_ir_tensor(expr.inputs[1]) + out_tensor = self.irgraph.get_tensor( + expr.outputs[0]._id, None, origin_tensor=inp_tensor + ) + qdtype = module.get_activation_dtype() + qparams = ( + module.act_fake_quant.get_qparams() + if hasattr(module.act_fake_quant, "get_qparams") + else module.act_observer.get_qparams() + ) + scale = qparams.scale + zero_point = qparams.zero_point + out_tensor.q_dtype = qdtype + out_tensor.scale = float(scale) + out_tensor.zero_point = int(zero_point) if zero_point else None + + else: + op_gen_cls = EXPR2OP.get(type(m.expr.value), None) + assert op_gen_cls, "Module {} is not supported.".format( + type(m.expr.value) + ) + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif isinstance(expr, CallFunction): + f = expr.func # expr.func.__module__ + "." + expr.func.__name__ + op_gen_cls = EXPR2OP.get(f, None) + assert op_gen_cls, "FUNCTION {} is not supported.".format(f) + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif isinstance(expr, Apply): + opdef = expr.opdef + op_gen_cls = EXPR2OP.get(str(opdef), None) + assert op_gen_cls, "OPDEF {} is not supported.".format(str(opdef)) + op = op_gen_cls(expr, self.irgraph).get_opr() + self.irgraph.add_op(op) + elif isinstance(expr, Input): + logger.warning("Do not suppot Input Expr.") + + def add_net_outputs(self): + for node in self.outputs: + assert ( + node._id in self.irgraph._tensor_ids + ), "output node is not generated by any opr" + out_tensor = self.tensor_resolver.get_ir_tensor(node, self) + self.irgraph.add_net_outputs(out_tensor) diff --git a/mgeconvert/frontend/tm_to_ir/tm_tensor_resolver.py b/mgeconvert/frontend/tm_to_ir/tm_tensor_resolver.py new file mode 100644 index 0000000..a96cd87 --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/tm_tensor_resolver.py @@ -0,0 +1,86 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module + +import megengine as mge +import numpy as np +from megengine.traced_module.expr import Constant, GetAttr +from megengine.traced_module.node import TensorNode + +from ...converter_ir.ir_graph import IRGraph +from ...converter_ir.ir_tensor import AxisOrder, IRTensor + + +class TensorNodeResolver: + __const_id = 0 + + def __init__(self, irgraph: IRGraph) -> None: + self.irgraph = irgraph + + def resolve(self, inp, owner_opr=None, param_name=None, axis_order=AxisOrder.NCHW): + scale = None + zero_point = None + + if isinstance(inp, TensorNode): + shape = inp.shape + dtype = inp.dtype + if isinstance(inp.expr, Constant): + np_data = inp.expr.value.numpy() + elif isinstance(inp.expr, GetAttr): + np_data = getattr(inp.expr.owner, inp.expr.name).numpy() + else: + np_data = None + name = inp._name + ori_id = inp._id + elif isinstance(inp, (int, float, list, np.ndarray)): + np_data = np.array(inp) + dtype = np_data.dtype.type + shape = np_data.shape + name = "const_val_" + str(TensorNodeResolver.__const_id) + TensorNodeResolver.__const_id += 1 + ori_id = None + elif isinstance(inp, mge.Tensor): + name = param_name + shape = inp.shape + dtype = inp.dtype + np_data = inp.numpy() + ori_id = None + + return ( + IRTensor( + name, + shape, + dtype, + scale=scale, + zero_point=zero_point, + np_data=np_data, + owner_opr=owner_opr, + axis=axis_order, + ), + ori_id, + ) + + def get_ir_tensor( + self, inp, owner_opr=None, user_opr=None, name=None, axis_order=AxisOrder.NCHW + ): + ir_tensor, ori_id = self.resolve( + inp, owner_opr=owner_opr, param_name=name, axis_order=axis_order + ) + ori_tensor = self.irgraph.get_tensor(ori_id, ir_tensor) + if user_opr is not None and user_opr not in ori_tensor.user_opr: + ori_tensor.add_user_opr(user_opr) + return ori_tensor + + def resolve_qparams(self, scale, zero_point): + if isinstance(scale, mge.Tensor): + scale = scale.numpy() + if zero_point: + if isinstance(scale, mge.Tensor): + zero_point = zero_point.numpy() + return scale, zero_point diff --git a/mgeconvert/frontend/tm_to_ir/tm_utils.py b/mgeconvert/frontend/tm_to_ir/tm_utils.py new file mode 100644 index 0000000..adfa70b --- /dev/null +++ b/mgeconvert/frontend/tm_to_ir/tm_utils.py @@ -0,0 +1,12 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from megengine import get_logger as mge_get_logger + + +def get_logger(*args): + return mge_get_logger(*args) diff --git a/mgeconvert/mge_context/__init__.py b/mgeconvert/mge_context/__init__.py deleted file mode 100644 index 711d7d9..0000000 --- a/mgeconvert/mge_context/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .mge_net import TopologyNetwork, get_platform, set_platform -from .mge_op import * -from .mge_tensor import Tensor -from .mge_transform import ( # type: ignore[attr-defined] - TransformerRule, - optimize_for_conversion, -) -from .mge_utils import get_dtype_name, get_logger diff --git a/mgeconvert/mge_context/mge_net.py b/mgeconvert/mge_context/mge_net.py deleted file mode 100644 index 5e72d02..0000000 --- a/mgeconvert/mge_context/mge_net.py +++ /dev/null @@ -1,144 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .mge_op import str_to_mge_class -from .mge_tensor import Tensor -from .mge_utils import ( - eval_partial, - get_dep_vars, - get_mge_version, - get_opr_type, - get_oprs_seq, - load_comp_graph_from_file, -) - - -class TopologyNetwork: - def __init__(self, model_path, outspec=None, prune_reshape=True): - _, outputs = load_comp_graph_from_file(model_path) - if outspec is not None: - output_spec = outspec.copy() - all_vars = get_dep_vars(outputs) + outputs - new_outputs = {} - for i in all_vars: - if i.name in output_spec: - new_outputs[i.name] = i - output_spec.remove(i.name) - assert len(output_spec) == 0, "Can not find {} in this model".format( - output_spec - ) - outputs = [new_outputs[i] for i in outspec] - all_oprs = get_oprs_seq(outputs, prune_reshape=prune_reshape) - self.input_vars = [] - self._orig_inputs = [] - self.output_vars = [] - self._orig_outputs = outputs - # why using two lists instead of one ordereddict for oprs and vars: - # for network trasformation, it's hard to manipulate ordereddict - # index of opr_ids and all_oprs should be matched - self._opr_ids = [] - self.all_oprs = [] - self._var_ids = [] - self.all_vars = [] - - for mge_opr in all_oprs: - if get_opr_type(mge_opr) == "Host2DeviceCopy": - self._orig_inputs.extend(mge_opr.outputs) - self.add_opr(mge_opr) - - for mge_opr in all_oprs: - opr = self.get_opr(mge_opr) - # set inp var - for x in mge_opr.inputs: - opr.add_inp_var(self.get_var(x)) - # set out var - for x in mge_opr.outputs: - opr.add_out_var(self.get_var(x)) - # set inp, out oprs - for x in mge_opr.inputs: - var = self.get_var(x) - if var.np_data is None: - owner = x.owner_opr if get_mge_version() <= "0.6.0" else x.owner - inp_opr = self.get_opr(owner) - if inp_opr is not None: - inp_opr.add_out_opr(opr) - opr.add_inp_opr(inp_opr) - - for x in self._orig_inputs: - self.input_vars.append(self.get_var(x)) - - for x in self._orig_outputs: - self.output_vars.append(self.get_var(x)) - - self.max_id = max([out_var.id for out_var in self.output_vars]) + 1 - # works iff all inputs have the same batch size - self.batch_size = self.input_vars[0].shape[0] - - def run(self, feed_input, end_op): - if end_op is None: - return eval_partial(feed_input, self.output_vars[-1].sym_var) - if isinstance(end_op, str): - for opr in self.all_oprs: - if opr.name == end_op: - end_op = opr - break - assert not isinstance(end_op, str), "end_op op does not exist" - if not len(end_op.out_vars): - return [] - flag = False - for var in end_op.out_vars: - if var.np_data is None: - flag = True - break - if not flag: - return [var.np_data for var in end_op.out_vars] - oup = [] - for var in end_op.out_vars: - oup.append(var.sym_var) - # WARN only calculate the first output var - oup = (oup[0],) - return eval_partial(feed_input, oup) - - def run_all(self, feed_input): - result_dict = {} - for opr in self.all_oprs: - out = self.run(feed_input, opr) - result_dict[opr.name] = out - return result_dict - - def add_opr(self, x): - assert x.id not in self._opr_ids - self._opr_ids.append(x.id) - self.all_oprs.append(str_to_mge_class(get_opr_type(x) + "Opr")(x)) - - def get_opr(self, x): - if x.id in self._opr_ids: - return self.all_oprs[self._opr_ids.index(x.id)] - else: - return None - - def get_var(self, x): - # auto convert to Tensor - if x.id not in self._var_ids: - owner = x.owner_opr if get_mge_version() <= "0.6.0" else x.owner - self._var_ids.append(x.id) - self.all_vars.append(Tensor(x, self.get_opr(owner))) - return self.all_vars[self._var_ids.index(x.id)] - - -class Config: - platform = "official" - - -def set_platform(platform): - assert platform in ["official", "mtk"] - Config.platform = platform - - -def get_platform(): - return Config.platform diff --git a/mgeconvert/mge_context/mge_op.py b/mgeconvert/mge_context/mge_op.py deleted file mode 100644 index 6dda4da..0000000 --- a/mgeconvert/mge_context/mge_op.py +++ /dev/null @@ -1,464 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import json -import sys -from typing import List # pylint: disable=unused-import - -from .mge_tensor import Tensor # pylint: disable=unused-import -from .mge_utils import ( - get_dep_vars, - get_mge_version, - get_opr_type, - get_shape, - get_symvar_value, -) - -mge_version = get_mge_version() - - -def str_to_mge_class(classname): - return getattr(sys.modules[__name__], classname) - - -class OpBase: - skip = False - name = "" - inp_vars = [] # type: List[Tensor] - out_vars = [] # type: List[Tensor] - - -class MgeOpr(OpBase): - def __init__(self, opr): - self._opr = opr - self.name = opr.name - self.name = self.name.replace(":", "_") - self.name = self.name.replace(".", "_") - self.name = self.name.replace(",", "_") - - self.id = opr.id - try: - self.type = get_opr_type(opr) - except AssertionError: - self.type = None - self.flag = None - self.inp_vars = [] - self.out_vars = [] - self.inp_oprs = [] - self.out_oprs = [] - self.params = opr.params if mge_version <= "0.6.0" else json.loads(opr.params) - self.activation = "IDENTITY" - - def add_inp_var(self, x): - self.inp_vars.append(x) - - def add_out_var(self, x): - if "workspace" in x.name: - return - self.out_vars.append(x) - - def add_inp_opr(self, x): - self.inp_oprs.append(x) - - def add_out_opr(self, x): - self.out_oprs.append(x) - - @property - def next_opr(self): - assert len(self.out_oprs) == 1 - return self.out_oprs[0] - - @property - def prev_opr(self): - assert len(self.inp_oprs) == 1 - return self.inp_oprs[0] - - -class PoolingForwardOpr(MgeOpr): - name = "PoolingForward" - - def __init__(self, opr): - super().__init__(opr) - self.data_format = self.params["format"] - self.mode = self.params["mode"] - - self.ph, self.pw = self.params["pad_h"], self.params["pad_w"] - self.sh, self.sw = self.params["stride_h"], self.params["stride_w"] - self.kh, self.kw = self.params["window_h"], self.params["window_w"] - - -class MatrixMulOpr(MgeOpr): - name = "MatrixMul" - - def __init__(self, opr): - super().__init__(opr) - self.format = self.params["format"] - self.transposeB = self.params["transposeB"] - self.transposeA = self.params["transposeA"] - self.compute_mode = self.params["compute_mode"] - - -class FullyConnectedOpr(MatrixMulOpr): - name = "FullyConnected" - - def __init__(self, name, mm_opr, param_bias): - super().__init__(mm_opr) - self.name = name - self.param_B = param_bias - - -class BatchedMatrixMulOpr(MgeOpr): - name = "BatchedMatrixMul" - - def __init__(self, opr): - super().__init__(opr) - self.format = self.params["format"] - self.transposeA = self.params["transposeA"] - self.transposeB = self.params["transposeB"] - self.compute_mode = self.params["compute_mode"] - - -class ConcatOpr(MgeOpr): - name = "Concat" - - def __init__(self, opr): - super().__init__(opr) - self.axis = self.params["axis"] - - -class ReshapeOpr(MgeOpr): - name = "Reshape" - - def __init__(self, opr): - super().__init__(opr) - self.input_shape = get_shape(opr.inputs[0]) - self.output_shape = get_shape(opr.outputs[0]) - self.shape_param = get_symvar_value(opr.inputs[1]) - - -class ConvolutionForwardOpr(MgeOpr): - name = "ConvolutionForward" - - def __init__(self, opr): - super().__init__(opr) - self.kernel_shape = get_shape(opr.inputs[1]) - if len(get_dep_vars(opr.inputs[1], "Host2DeviceCopy")) == 0: - self.param_W = get_symvar_value(opr.inputs[1]) - else: - self.param_W = None - self.data_format = self.params["format"] - self.dilation_w = self.params["dilate_w"] - self.dilation_h = self.params["dilate_h"] - self.compute_mode = self.params["compute_mode"] - self.sparse = self.params["sparse"] - - self.ph, self.pw = self.params["pad_h"], self.params["pad_w"] - self.sh, self.sw = self.params["stride_h"], self.params["stride_w"] - if self.data_format == "NCHW": - self.kh = self.kernel_shape[-2] - self.kw = self.kernel_shape[-1] - else: - assert False, "do not support this {} format".format(self.data_format) - - self.num_output = get_shape(opr.outputs[0])[1] - self.bias_term = False - self.group = self.kernel_shape[0] if len(self.kernel_shape) == 5 else 1 - - -class ConvBiasForwardOpr(ConvolutionForwardOpr): - name = "ConvBiasForward" - - def __init__(self, opr): - super().__init__(opr) - self.activation = self.params["nonlineMode"] - - -class ConvForwardBiasOpr(ConvolutionForwardOpr): - name = "ConvForwardBias" - - def __init__(self, name, conv_opr, param_bias): - super().__init__(conv_opr) - self.name = name - self.bias_term = True - self.param_B = param_bias - - -class ElemwiseOpr(MgeOpr): - name = "Elemwise" - - def __init__(self, opr): - super().__init__(opr) - try: - self.mode = self.params["mode"] - if "RELU" in self.mode: - self.activation = "RELU" - except RuntimeError: - self.mode = "NONE" - - -class ElemwiseMultiTypeOpr(MgeOpr): - name = "ElemwiseMultiType" - - -class Host2DeviceCopyOpr(MgeOpr): - name = "Host2DeviceCopy" - - def __init__(self, opr): - super().__init__(opr) - assert len(opr.outputs) == 1, "wrong number of outputs" - self.shape = get_shape(opr.outputs[0]) - - -class MultipleDeviceTensorHolderOpr(MgeOpr): - name = "MultipleDeviceTensorHolder" - - -class VolatileSharedDeviceTensorOpr(MgeOpr): - name = "VolatileSharedDeviceTensor" - - -class IndexingOneHotOpr(MgeOpr): - name = "IndexingOneHotOpr" - - -class SubtensorOpr(MgeOpr): - name = "Subtensor" - - def __init__(self, opr): - super().__init__(opr) - self.has_step = [] - self.has_begin = [] - self.has_end = [] - self.has_idx = [] - self.axis = [] - for param in self.params: - self.has_step.append(param["step"]) - self.has_begin.append(param["begin"]) - self.has_end.append(param["end"]) - self.has_idx.append(param["idx"]) - self.axis.append(param["axis"]) - - begin_param = [] - end_param = [] - step_param = [] - squeeze_axis = [] - slice_param = [get_symvar_value(v)[0] for v in opr.inputs[1:]] - for i in range(len(self.has_begin)): - if self.has_idx[i] == 1: - begin_idx = slice_param.pop(0) - end_idx = begin_idx + 1 - begin_param.append(begin_idx) - end_param.append(end_idx) - step_param.append(1) - squeeze_axis.append(self.axis[i]) - else: - if self.has_begin[i]: - begin_param.append(slice_param.pop(0)) - else: - begin_param.append(0) - if self.has_end[i]: - end_param.append(slice_param.pop(0)) - else: - end_param.append(2147483647) - step_param.append(1 if self.has_step[i] == 0 else slice_param.pop(0)) - - self.squeeze_axis = squeeze_axis - self.begin_param = begin_param - self.end_param = end_param - self.step_param = step_param - - -class ImmutableTensorOpr(MgeOpr): - name = "ImmutableTensor" - - -class GetVarShapeOpr(MgeOpr): - name = "GetVarShape" - - -class BatchNormForwardOpr(MgeOpr): - name = "BatchNormForward" - - def __init__(self, opr): - super().__init__(opr) - self.output_idx = 4 if mge_version <= "0.6.0" else 2 - self.scale = get_symvar_value(opr.inputs[1]).squeeze() - self.bias = get_symvar_value(opr.inputs[2]).squeeze() - self.mean = get_symvar_value(opr.inputs[3]).squeeze() - self.var = get_symvar_value(opr.inputs[4]).squeeze() - - -class MarkNoBroadcastElemwiseOpr(MgeOpr): - name = "MarkNoBroadcastElemwise" - - def __init__(self, opr): - super().__init__(opr) - self.mode = "Identity" - - -class IdentityOpr(MgeOpr): - name = "Identity" - - def __init__(self, opr): - super().__init__(opr) - self.mode = "Identity" - - -class SharedDeviceTensorOpr(MgeOpr): - name = "SharedDeviceTensor" - - def __init__(self, opr): - super().__init__(opr) - assert len(opr.outputs) == 1, "wrong number of outputs" - self.shape = get_shape(opr.outputs[0]) - - -class DimshuffleOpr(MgeOpr): - name = "Dimshuffle" - - def __init__(self, opr): - super().__init__(opr) - self.pattern = self.params["pattern"] - self.ndim = self.params["ndim"] - - -class TypeCvtOpr(MgeOpr): - name = "TypeCvt" - - -class ReduceOpr(MgeOpr): - name = "Reduce" - - def __init__(self, opr): - super().__init__(opr) - self.axis = self.params["axis"] - self.mode = self.params["mode"] - - -class AxisAddRemoveOpr(MgeOpr): - name = "AxisAddRemove" - - def __init__(self, opr): - super().__init__(opr) - self.desc = self.params["desc"] - self.nr_desc = self.params["nr_desc"] - self.output_shape = get_shape(opr.outputs[0]) - - -class BroadcastOpr(MgeOpr): - name = "Broadcast" - - -class WarpPerspectiveForwardOpr(MgeOpr): - name = "WarpPerspectiveForward" - - -class LinspaceOpr(MgeOpr): - name = "Linspace" - - -class ConvolutionBackwardDataOpr(MgeOpr): - name = "ConvolutionBackwardData" - - def __init__(self, opr): - super().__init__(opr) - # opr.inputs[0] is conv kernel - self.kernel_shape = get_shape(opr.inputs[0]) - self.param_W = get_symvar_value(opr.inputs[0]) - self.data_format = self.params["format"] - self.dilation_w = self.params["dilate_w"] - self.dilation_h = self.params["dilate_h"] - self.compute_mode = self.params["compute_mode"] - self.sparse = self.params["sparse"] - - self.ph, self.pw = self.params["pad_h"], self.params["pad_w"] - self.sh, self.sw = self.params["stride_h"], self.params["stride_w"] - if self.data_format == "NCHW": - self.kh = self.kernel_shape[-2] - self.kw = self.kernel_shape[-1] - else: - assert False, "do not support this {} format".format(self.data_format) - - self.num_output = get_shape(opr.outputs[0])[1] - self.bias_term = False - self.group = self.param_W.shape[0] if self.param_W.ndim == 5 else 1 - - -class ConvolutionBackwardFilterOpr(MgeOpr): - r"""refer to https://github.com/pytorch/pytorch/blob/master/torch/nn/grad.py#L170 - """ - - name = "ConvolutionBackwardFilter" - - def __init__(self, opr): - super().__init__(opr) - assert self.params["format"] == "NCHW", "do not support this {}".format( - self.params["format"] - ) - - src, grad_out, weight = opr.inputs - - self.ni, self.ci, self.hi, self.wi = src.shape - self.no, self.co, self.ho, self.wo = grad_out.shape - assert self.ni == self.no - - self.dilate_w = self.params["dilate_w"] - self.dilate_h = self.params["dilate_h"] - self.ph, self.pw = self.params["pad_h"], self.params["pad_w"] - self.sh, self.sw = self.params["stride_h"], self.params["stride_w"] - - if len(weight.shape) == 5: - self.group = weight.shape[0] - self.kh = weight.shape[3] - self.kw = weight.shape[4] - else: - self.group = 1 - self.kh = weight.shape[2] - self.kw = weight.shape[3] - - -class ConvolutionBackwardDataBiasOpr(ConvolutionBackwardDataOpr): - name = "ConvolutionBackwardDataBias" - - def __init__(self, name, opr, param_bias): - super().__init__(opr) - self.name = name - self.param_B = param_bias - self.bias_term = True - # opr.inputs[0] is conv kernel - - -class ResizeForwardOpr(MgeOpr): - name = "ResizeForward" - - -class LeakyReluOpr(MgeOpr): - name = "LeakyRelu" - - def __init__(self, name, opr, negative_slope): - super().__init__(opr) - self.name = name - self.negative_slope = negative_slope - - -class Relu6Opr(OpBase): - name = "Relu6" - - -class SoftmaxOpr(OpBase): - name = "Softmax" - beta = 0 - - -class PadOpr(OpBase): - name = "Pad" - - -class SqueezeOpr(OpBase): - name = "Squeeze" - squeeze_dims = [] # type: ignore[var-annotated] diff --git a/mgeconvert/mge_context/mge_tensor.py b/mgeconvert/mge_context/mge_tensor.py deleted file mode 100644 index 9e52094..0000000 --- a/mgeconvert/mge_context/mge_tensor.py +++ /dev/null @@ -1,68 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .mge_utils import get_dep_vars, get_shape, get_symvar_value - - -class FakeSymbolVar: - def __init__(self, sid, name, shape, dtype, owner=None, byte_list=None): - self.id = sid - self.name = name - self.shape = shape - self.dtype = dtype - self.owner = owner - self.byte_list = byte_list - - # in order to be compatible with v0.6.0 - @property - def owner_opr(self): - return self.owner - - def _get_imm_shape(self): - return self.shape - - -class Tensor: - def __init__(self, sym_var, owner_opr): - self.is_faked = isinstance(sym_var, FakeSymbolVar) - self.id = sym_var.id - self._var = sym_var - self.name = sym_var.name - self.name = self.name.replace(":", "_") - self.name = self.name.replace(".", "_") - self.name = self.name.replace(",", "_") - self.shape = get_shape(sym_var) - self.owner_opr = owner_opr - if self.shape is not None: - self.ndim = len(self.shape) - try: - self.dtype = sym_var.dtype - except: # pylint: disable=bare-except - self.dtype = None - - self.byte_list = getattr(sym_var, "byte_list", None) - - self.np_data = None - try: - if len(get_dep_vars(sym_var, "Host2DeviceCopy")) == 0: - self.np_data = get_symvar_value(sym_var) - except: # pylint: disable=bare-except - self.np_data = None - - try: - self.qbit = self.dtype.metadata["mgb_dtype"]["name"] - except: # pylint: disable=bare-except - self.qbit = None - - try: - self.scale = self.dtype.metadata["mgb_dtype"]["scale"] - except: # pylint: disable=bare-except - self.scale = None - - def set_owner_opr(self, owner_opr): - self.owner_opr = owner_opr diff --git a/mgeconvert/mge_context/mge_transform.py b/mgeconvert/mge_context/mge_transform.py deleted file mode 100644 index 8ea6b00..0000000 --- a/mgeconvert/mge_context/mge_transform.py +++ /dev/null @@ -1,959 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - -# type: ignore[no-redef] - -from collections import OrderedDict -from enum import Enum -from typing import Callable, Dict, List, Sequence - -import numpy as np - -from . import mge_op as Ops -from .mge_op import ( - ConvBiasForwardOpr, - ConvForwardBiasOpr, - ConvolutionBackwardDataOpr, - ConvolutionForwardOpr, - DimshuffleOpr, - ElemwiseOpr, - OpBase, - PadOpr, - PoolingForwardOpr, - ReduceOpr, - Relu6Opr, - ReshapeOpr, - SoftmaxOpr, - TypeCvtOpr, -) -from .mge_tensor import FakeSymbolVar, Tensor -from .mge_utils import isconst - - -class PatternNode: - def __init__(self, type, is_output=False, const_value=None): - self.op = None - self.type = type - self.inp_oprs = [] - self.inp_const = [] - self.inp_vars = [] - self.is_output = is_output - self.const_value = const_value - - def check_const_value(self, op): - inp_vars = [v.np_data for v in op.inp_vars] - for const in self.const_value: - idx = const[0] - if idx == -1: - find = False - for index, v in enumerate(inp_vars): - if np.array_equal(const[1], v): - find = True - del inp_vars[index] - break - if not find: - return False - elif not np.array_equal(const[1], inp_vars[idx]): - return False - return True - - -get_type = lambda op: op.mode if isinstance(op, Ops.ElemwiseOpr) else type(op).__name__ - - -def match(node, opr): - node_queue = [node] - opr_queue = [opr] - matched_opr = set() - matched_node = set() - while len(node_queue) != 0: - cur_node = node_queue.pop(0) - cur_opr = opr_queue.pop(0) - if cur_node.type != get_type(cur_opr) and cur_node.type != "*" or cur_opr.skip: - return False - if cur_node.op == None: - cur_node.op = cur_opr - if cur_node.const_value != None: - if not cur_node.check_const_value(cur_opr): - return False - elif cur_node.op != cur_opr: - return False - - matched_opr.add(cur_opr) - matched_node.add(cur_node) - for i, var in enumerate(cur_opr.inp_vars): - if isconst(var): - cur_node.inp_const.append([i, var.np_data]) - else: - cur_node.inp_vars.append([i, var]) - if len(cur_node.inp_oprs) == 0: - continue - if len(cur_node.inp_oprs) != len(cur_opr.inp_oprs): - return False - - for i, j in zip(cur_node.inp_oprs, cur_opr.inp_oprs): - node_queue.append(i) - opr_queue.append(j) - - for n in matched_node: - if n.is_output: - continue - for op in n.op.out_oprs: - if op not in matched_opr: - return False - - return True - - -class TransformerRule(Enum): - # general rules - NOPE = 1 - - # for TFLite - REDUCE_AXIS_AS_INPUT = 100 - REMOVE_RESHAPE_INPUT = 101 - # FUSE_FOR_RELU6 pass should happen before FUSE_ACTIVATION - FUSE_FOR_RELU6 = 102 - FUSE_ACTIVATION = 103 - CONV_ADD_ZERO_BIAS = 104 - DEPTHWISE_CONV_RESHAPE_WEIGHT = 105 - FUSE_SOFTMAX = 106 - DECONV_SHAPE_AS_INPUT = 107 - FUSE_ASTYPE = 108 - PADDING_FOR_CONV = 109 - TRANSPOSE_PATTERN_AS_INPUT = 110 - # FUSE_FOR_LEAKY_RELU should happen before EXPAND_MUL_ADD3 - FUSE_FOR_LEAKY_RELU = 111 - EXPAND_MUL_ADD3 = 112 - EXPAND_ADD_SIGMOID = 113 - FUSE_FOR_CONV_BIAS = 114 - FUSE_FOR_DECONV_BIAS = 115 - FUSE_FOR_FULLY_CONNECTED = 116 - RESHAPE_BIAS_TO_1DIM = 117 - # for TFLite Converter - SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE = 200 - - -TRANSFORMMAP: Dict[Enum, Callable] = {} - - -def optimize_for_conversion(net, transformer_options): - if not isinstance(transformer_options, Sequence): - transformer_options = (transformer_options,) - for option in transformer_options: - TRANSFORMMAP[option](net) - - -def _register_tranformation_rule(transformer_option): - def callback(impl): - TRANSFORMMAP[transformer_option] = impl - - return callback - - -@_register_tranformation_rule(TransformerRule.REMOVE_RESHAPE_INPUT) -def _remove_reshape_input(net): - for op in net.all_oprs: - if not isinstance(op, ReshapeOpr): - continue - - del op.inp_vars[1] - - -@_register_tranformation_rule(TransformerRule.REDUCE_AXIS_AS_INPUT) -def _reduce_axis_as_input(net): - for op in net.all_oprs: - if not isinstance(op, ReduceOpr): - continue - - byte_list = np.int32(op.axis).tobytes() - - axis_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_axis", - shape=[1], - dtype=np.int32, - owner=op, - byte_list=byte_list, - ) - net.max_id += 1 - - axis_tensor = net.get_var(axis_symvar) - op.add_inp_var(axis_tensor) - - -@_register_tranformation_rule(TransformerRule.FUSE_FOR_RELU6) -def _fuse_for_relu6(net): - matches = OrderedDict() - - for op in net.all_oprs: - if not isinstance(op, ElemwiseOpr): - continue - if len(op.inp_oprs) <= 0 or len(op.inp_oprs) >= 2: - continue - prev_op = op.inp_oprs[0] - cur_index = net._opr_ids.index(op.id) - - if op.mode == "MIN" and np.array_equal(op.inp_vars[1].np_data, np.array([6])): - if ( - isinstance(prev_op, ElemwiseOpr) - and prev_op.mode == "MAX" - and np.array_equal(prev_op.inp_vars[1].np_data, np.array([0])) - and net._opr_ids.index(prev_op.id) == cur_index - 1 - ): - relu6_opr = Relu6Opr() - relu6_opr.inp_vars = prev_op.inp_vars - relu6_opr.out_vars = op.out_vars - relu6_opr.inp_oprs = [prev_op.prev_opr] - relu6_opr.out_oprs = op.out_oprs - matches[prev_op.id] = (net.max_id, relu6_opr) - net.max_id += 1 - if op.mode == "MAX" and np.array_equal(op.inp_vars[1].np_data, np.array([0])): - if ( - isinstance(prev_op, ElemwiseOpr) - and prev_op.mode == "MIN" - and np.array_equal(prev_op.inp_vars[1].np_data, np.array([6])) - and net._opr_ids.index(prev_op.id) == cur_index - 1 - ): - relu6_opr = Relu6Opr() - relu6_opr.inp_vars = prev_op.inp_vars - relu6_opr.out_vars = op.out_vars - relu6_opr.inp_oprs = [prev_op.inp_oprs[0]] - relu6_opr.out_oprs = op.out_oprs - matches[prev_op.id] = (net.max_id, relu6_opr) - net.max_id += 1 - - for original_id, generated_pair in list(matches.items())[::-1]: - index = net._opr_ids.index(original_id) - del net._opr_ids[index : index + 2] - del net.all_oprs[index : index + 2] - - net._opr_ids.insert(index, generated_pair[0]) - net.all_oprs.insert(index, generated_pair[1]) - - -@_register_tranformation_rule(TransformerRule.FUSE_ACTIVATION) -def _fuse_activation(net): - delete_intended = [] - - for op_id, op in zip(net._opr_ids, net.all_oprs): - if isinstance(op, Relu6Opr) or ( - isinstance(op, ElemwiseOpr) and op.mode in ("RELU", "TANH") - ): - prev_op = op.inp_oprs[0] - if prev_op.activation != "IDENTITY": - continue - - # activation(relu/relu6/tanh) must be fused with previous opr - activation = getattr(op, "mode", "IDENTITY") - activation = "RELU6" if isinstance(op, Relu6Opr) else activation - prev_op.activation = activation - prev_op.out_vars = op.out_vars - - for post_op in op.out_oprs: - idx = post_op.inp_oprs.index(op) - post_op.inp_oprs[idx] = prev_op - if post_op not in prev_op.out_oprs: - prev_op.out_oprs.append(post_op) - - delete_intended.append(net._opr_ids.index(op_id)) - - for delete_idx in delete_intended[::-1]: - del net._opr_ids[delete_idx] - del net.all_oprs[delete_idx] - - -@_register_tranformation_rule(TransformerRule.CONV_ADD_ZERO_BIAS) -def _conv_add_zero_bias(net): - for op in net.all_oprs: - if not isinstance(op, ConvolutionForwardOpr): - continue - if isinstance(op, (ConvBiasForwardOpr, ConvForwardBiasOpr)): - continue - - result_shape = [op.out_vars[0].shape[1]] - if op.inp_vars[0].dtype == np.float32: - dtype = np.float32 - else: # op.inp_vars[0].dtype == np.uint8 - scale0 = op.inp_vars[0].dtype.metadata["mgb_dtype"]["scale"] - scale1 = op.inp_vars[1].dtype.metadata["mgb_dtype"]["scale"] - dtype = np.dtype( - "int32", - metadata={ - "mgb_dtype": {"name": "Quantized8Asymm", "scale": scale0 * scale1} - }, - ) - - byte_list = np.zeros( - result_shape, np.int32 if dtype == np.int32 else np.float32 - ).tobytes() - bias_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_bias", - shape=result_shape, - dtype=dtype, - owner=op, - byte_list=byte_list, - ) - net.max_id += 1 - - bias_tensor = net.get_var(bias_symvar) - op.add_inp_var(bias_tensor) - - -@_register_tranformation_rule(TransformerRule.DEPTHWISE_CONV_RESHAPE_WEIGHT) -def _depthwise_conv_reshape_weight(net): - # general group conv is not supported for TFLite - for op in net.all_oprs: - if not isinstance(op, ConvolutionForwardOpr): - continue - if op.group == 1: - continue - - var = op.inp_vars[1] - ic, cm = var.shape[1], var.shape[2] * op.group - h, w = var.shape[3:5] - var.shape = (ic, cm, h, w) - var.ndim = len(var.shape) - var.np_data.reshape(ic, cm, h, w) - - -@_register_tranformation_rule(TransformerRule.FUSE_SOFTMAX) -def _fuse_softmax(net): - matches = OrderedDict() - - for op in net.all_oprs: - if not isinstance(op, ElemwiseOpr) or op.mode != "TRUE_DIV": - continue - try: - prev_op = op.inp_oprs[1] - cur_index = net._opr_ids.index(op.id) - if ( - not isinstance(prev_op, ReduceOpr) - or prev_op.mode != "SUM" - or prev_op.axis != 1 - or net._opr_ids.index(prev_op.id) != cur_index - 1 - ): - continue - prev_op = op.inp_oprs[0] - if ( - not isinstance(prev_op, ElemwiseOpr) - or prev_op.mode != "EXP" - or net._opr_ids.index(prev_op.id) != cur_index - 2 - ): - continue - prev_op = prev_op.prev_opr - if ( - not isinstance(prev_op, ElemwiseOpr) - or prev_op.mode != "SUB" - or net._opr_ids.index(prev_op.id) != cur_index - 3 - ): - continue - prev_op = prev_op.inp_oprs[1] - if ( - not isinstance(prev_op, ReduceOpr) - or prev_op.mode != "MAX" - or prev_op.axis != 1 - or net._opr_ids.index(prev_op.id) != cur_index - 4 - ): - continue - except IndexError: # doesn't match - continue - - softmax_opr = SoftmaxOpr() - softmax_opr.beta = 1 - softmax_opr.inp_vars = prev_op.inp_vars[:1] - softmax_opr.out_vars = op.out_vars - softmax_opr.inp_oprs = [prev_op.inp_oprs[0]] - softmax_opr.out_oprs = op.out_oprs - softmax_opr.prev_opr = prev_op.inp_oprs[0] - matches[prev_op.id] = (net.max_id, softmax_opr) - net.max_id += 1 - - for original_id, generated_pair in list(matches.items())[::-1]: - index = net._opr_ids.index(original_id) - del net._opr_ids[index : index + 5] - del net.all_oprs[index : index + 5] - - net._opr_ids.insert(index, generated_pair[0]) - net.all_oprs.insert(index, generated_pair[1]) - - -@_register_tranformation_rule(TransformerRule.DECONV_SHAPE_AS_INPUT) -def _deconv_shape_as_input(net): - for op in net.all_oprs: - if not isinstance(op, ConvolutionBackwardDataOpr): - continue - - byte_list = [] - result_shape = op.out_vars[0].shape - number_list = [ - result_shape[0], - result_shape[2], - result_shape[3], - result_shape[1], - ] - for i in number_list: - byte_list.extend(np.int32(i).tobytes()) - shape_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_shape", - shape=[4], - dtype=np.int32, - owner=op, - byte_list=byte_list, - ) - net.max_id += 1 - shape_tensor = net.get_var(shape_symvar) - op.inp_vars = [shape_tensor, op.inp_vars[1], op.inp_vars[0]] - - -@_register_tranformation_rule(TransformerRule.SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE) -def _make_slice_as_inputs(net): - for op in net.all_oprs: - if not isinstance(op, Ops.SubtensorOpr): - continue - - ndim = op.inp_vars[0].ndim - - def make_input(axis, param, init_value): - # make inputs: begin, end and step. - ret = [init_value] * ndim # pylint: disable=cell-var-from-loop - for k, v in zip(axis, param): - ret[k] = v - ret = FakeSymbolVar( - sid=net.max_id, - name=op.name + "fake_input", # pylint: disable=cell-var-from-loop - shape=[len(ret)], - dtype=np.int32, - owner=op, # pylint: disable=cell-var-from-loop - byte_list=np.array(ret, np.int32).tobytes(), - ) - net.max_id += 1 - return net.get_var(ret) - - begins_tensor = make_input(op.axis, op.begin_param, 0) - ends_tensor = make_input(op.axis, op.end_param, np.iinfo(np.int32).max) - steps_tensor = make_input(op.axis, op.step_param, 1) - - op.inp_vars = [op.inp_vars[0], begins_tensor, ends_tensor, steps_tensor] - - # TFLite slice do not support squeeze axis, so insert a squeeze opr here. - # infer actual output shape of tflite slice - desired_out_shape = op.out_vars[0].shape - actual_out_shape = [1] * ndim - idx = 0 - for i in range(ndim): - if i in op.squeeze_axis: - continue - actual_out_shape[i] = desired_out_shape[idx] - idx += 1 - slice_out_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "fake_output", - shape=actual_out_shape, - dtype=op.out_vars[0].dtype, - owner=op, - ) - net.max_id += 1 - slice_op_output = net.get_var(slice_out_symvar) - old_out = op.out_vars - op.out_vars = [slice_op_output] - - squeeze = Ops.SqueezeOpr() - squeeze.squeeze_dims = op.squeeze_axis - squeeze.inp_vars = [slice_op_output] - squeeze.out_vars = old_out - squeeze.inp_oprs = [op] - squeeze.out_oprs = op.out_oprs - op.out_oprs = [squeeze] - squeeze.id = net.max_id - net.max_id += 1 - - idx = net._opr_ids.index(op.id) + 1 - net._opr_ids.insert(idx, squeeze.id) - net.all_oprs.insert(idx, squeeze) - - -@_register_tranformation_rule(TransformerRule.PADDING_FOR_CONV) -def _make_padding(net): - def have_padding(opr): - if ( - hasattr(opr, "ph") - and (opr.ph > 0 or opr.pw > 0) - and not isinstance(opr, ConvolutionBackwardDataOpr) - ): - return True - return False - - insert_intended = OrderedDict() - - for op in net.all_oprs: - if type(op) not in ( # pylint: disable=unidiomatic-typecheck - ConvolutionForwardOpr, - ConvBiasForwardOpr, - PoolingForwardOpr, - ): - continue - - if have_padding(op): - assert op.inp_vars[0].ndim == 4, "ERROR: unsupported padding mode" - byte_list = [] - number_list = [0, 0, op.ph, op.ph, op.pw, op.pw, 0, 0] - for i in number_list: - byte_list.extend(np.int32(i).tobytes()) - pad_in_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_pad_in", - shape=[4, 2], - dtype=np.int32, - owner=None, - byte_list=byte_list, - ) - net.max_id += 1 - net._var_ids.append(pad_in_symvar.id) - pad_in_tensor = Tensor(pad_in_symvar, None) - net.all_vars.append(pad_in_tensor) - - shape = list(op.inp_vars[0].shape) - pad_out_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_pad_out", - shape=[shape[0], shape[2] + op.ph * 2, shape[3] + op.pw * 2, shape[1],], - dtype=op.inp_vars[0].dtype, - owner=None, - byte_list=None, - ) - net.max_id += 1 - net._var_ids.append(pad_out_symvar.id) - pad_out_tensor = Tensor(pad_out_symvar, None) - net.all_vars.append(pad_out_tensor) - - pad_opr = PadOpr() - pad_opr.type = "PadOpr" - pad_opr.name = "pad_" + op.name - pad_opr.id = net.max_id - net.max_id += 1 - pad_opr.inp_vars = [op.inp_vars[0], pad_in_tensor] - pad_opr.out_vars = [pad_out_tensor] - pad_opr.inp_oprs = [op.inp_oprs[0]] - pad_opr.out_oprs = [op] - pad_out_tensor.owner = pad_opr - op.inp_vars = [pad_out_tensor] + op.inp_vars[1:] - op.inp_oprs = [pad_opr] + op.inp_oprs[1:] - - index = net._opr_ids.index(op.id) - insert_intended[index] = (pad_opr.id, pad_opr) - - for index, generated_pair in list(insert_intended.items())[::-1]: - net._opr_ids.insert(index, generated_pair[0]) - net.all_oprs.insert(index, generated_pair[1]) - - -@_register_tranformation_rule(TransformerRule.FUSE_ASTYPE) -def _fuse_astype(net): - def check_dtype(opr, dtype1, dtype2): - if opr.inp_vars[0].dtype == dtype1 and opr.out_vars[0].dtype == dtype2: - return True - return False - - delete_intended = [] - - for op_id, op in zip(net._opr_ids, net.all_oprs): - if isinstance(op, TypeCvtOpr): - # typecvt.prev_opr must have single output - prev_op = op.prev_opr - if ( - check_dtype(op, np.float32, np.int32) - or check_dtype(op, np.float32, np.uint8) - or check_dtype(op, np.float32, np.int8) - or check_dtype(op, np.int32, np.int8) - ): # quant phase - is_net_input = prev_op.out_vars[0] in net.input_vars - if is_net_input: - net.input_vars.remove(prev_op.out_vars[0]) - prev_op.out_vars[0] = op.out_vars[0] - op.out_vars[0].owner = prev_op - if is_net_input: - net.input_vars.append(prev_op.out_vars[0]) - delete_intended.append(net._opr_ids.index(op_id)) - else: # dequant phase, typecvt must be the last opr in the model - if ( - check_dtype(op, np.int8, np.int32) - or check_dtype(op, np.uint8, np.float32) - or check_dtype(op, np.int8, np.float32) - ): - is_net_output = op.out_vars[0] in net.output_vars - if is_net_output: - net.output_vars.remove(op.out_vars[0]) - net.output_vars.append(prev_op.out_vars[0]) - delete_intended.append(net._opr_ids.index(op_id)) - - for delete_idx in delete_intended[::-1]: - del net._opr_ids[delete_idx] - del net.all_oprs[delete_idx] - - -@_register_tranformation_rule(TransformerRule.TRANSPOSE_PATTERN_AS_INPUT) -def _transpose_pattern_as_input(net): - for op in net.all_oprs: - if not isinstance(op, DimshuffleOpr): - continue - - pat = op.pattern - assert len(pat) == 4, "only 4D input is supported by Dimshuffle" - # NCHW perm -> NHWC perm, (0, 3, 1, 2) is inv(0, 2, 3, 1) - pat = [pat[0], pat[3], pat[1], pat[2]] - byte_list = [] - for i in pat: - byte_list.extend(np.int32(i).tobytes()) - - pattern_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_pattern", - shape=[len(pat)], - dtype=np.int32, - owner=op, - byte_list=byte_list, - ) - net.max_id += 1 - - pattern_tensor = net.get_var(pattern_symvar) - op.add_inp_var(pattern_tensor) - - -@_register_tranformation_rule(TransformerRule.EXPAND_MUL_ADD3) -def _expand_mul_add3(net): - insert_intended = OrderedDict() - - for op in net.all_oprs: - if not isinstance(op, ElemwiseOpr): - continue - if op.mode != "FUSE_MUL_ADD3": - continue - - mul_out_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_mul_out", - shape=op.inp_vars[0].shape, - dtype=op.inp_vars[0].dtype, - owner=None, - byte_list=None, - ) - net.max_id += 1 - net._var_ids.append(mul_out_symvar.id) - mul_out_tensor = Tensor(mul_out_symvar, None) - net.all_vars.append(mul_out_tensor) - - mul_config = OpBase() - mul_config.name = "mul_" + op.name - mul_config.id = net.max_id - mul_config.params = '{"mode": "MUL"}' - mul_opr = ElemwiseOpr(mul_config) - mul_opr.type = "ElemwiseOpr" - net.max_id += 1 - mul_out_tensor.owner = mul_opr - mul_opr.inp_vars = op.inp_vars[:2] - mul_opr.out_vars = [mul_out_tensor] - mul_opr.inp_oprs = op.inp_oprs[:2] - mul_opr.out_oprs = [op] - - op.mode = "ADD" - op.inp_vars = [mul_out_tensor, op.inp_vars[2]] - op.inp_oprs = [mul_opr] - if len(op.inp_oprs) > 2: - op.inp_oprs.append(op.inp_oprs[2]) - - index = net._opr_ids.index(op.id) - insert_intended[index] = (mul_opr.id, mul_opr) - - for index, generated_pair in list(insert_intended.items())[::-1]: - net._opr_ids.insert(index, generated_pair[0]) - net.all_oprs.insert(index, generated_pair[1]) - - -@_register_tranformation_rule(TransformerRule.EXPAND_ADD_SIGMOID) -def _expand_add_sigmoid(net): - insert_intended = OrderedDict() - - for op in net.all_oprs: - if not isinstance(op, ElemwiseOpr): - continue - if op.mode != "FUSE_ADD_SIGMOID": - continue - - add_out_symvar = FakeSymbolVar( - sid=net.max_id, - name=op.name + "_add_out", - shape=op.inp_vars[0].shape, - dtype=op.inp_vars[0].dtype, - owner=None, - byte_list=None, - ) - net.max_id += 1 - net._var_ids.append(add_out_symvar.id) - add_out_tensor = Tensor(add_out_symvar, None) - net.all_vars.append(add_out_tensor) - - sigmoid_config = OpBase() - sigmoid_config.name = "sigmoid_" + op.name - sigmoid_config.id = net.max_id - sigmoid_config.params = '{"mode": "SIGMOID"}' - sigmoid_opr = ElemwiseOpr(sigmoid_config) - sigmoid_opr.type = "ElemwiseOpr" - net.max_id += 1 - - sigmoid_opr.inp_vars = [add_out_tensor] - sigmoid_opr.out_vars = op.out_vars - sigmoid_opr.inp_oprs = [op] - sigmoid_opr.out_oprs = op.out_oprs - - add_out_tensor.owner = op - op.mode = "ADD" - op.out_vars = [add_out_tensor] - op.out_oprs = [sigmoid_opr] - - index = net._opr_ids.index(op.id) + 1 - insert_intended[index] = (sigmoid_opr.id, sigmoid_opr) - - for index, generated_pair in list(insert_intended.items())[::-1]: - net._opr_ids.insert(index, generated_pair[0]) - net.all_oprs.insert(index, generated_pair[1]) - - -def _fuse_for_leaky_relu(opr): - assert ( - len(opr.out_oprs) == 1 - and isinstance(opr.out_oprs[0], Ops.ElemwiseOpr) - and opr.out_oprs[0].mode == "ADD" - ) - add_node = PatternNode("ADD", is_output=True) - mul_node = PatternNode("MUL") - max_node = PatternNode("MAX", const_value=[(-1, [0.0])]) - min_node = PatternNode("MIN", const_value=[(-1, [0.0])]) - add_node.inp_oprs = [max_node, mul_node] - mul_node.inp_oprs = [min_node] - - add_opr = opr.out_oprs[0] - if match(add_node, add_opr): - if ( - max_node.inp_vars[0] == min_node.inp_vars[0] - and len(mul_node.inp_const) == 1 - and mul_node.inp_const[0][1].shape == (1,) - ): - leaky_relu = Ops.LeakyReluOpr( - "leaky_relu_" + add_node.op.name, - add_node.op._opr, - mul_node.inp_const[0][1], - ) - leaky_relu.inp_vars = [max_node.inp_vars[0][1]] - leaky_relu.out_vars = add_node.op.out_vars - leaky_relu.inp_oprs = max_node.op.inp_oprs - leaky_relu.out_oprs = add_node.op.out_oprs - for node in [add_node, mul_node, max_node, min_node]: - node.op.skip = True - return leaky_relu - - return None - - -@_register_tranformation_rule(TransformerRule.FUSE_FOR_LEAKY_RELU) -def _(net): - """ - Elemwise(ADD) + Elemwise(MUL) + Elemwise(MAX) + Elemwise(MIN) -> LeakyRelu - """ - matches = list() - for opr in net.all_oprs: - if ( - get_type(opr) == "MAX" - and len(opr.out_oprs) == 1 - and get_type(opr.out_oprs[0]) == "ADD" - ): - leaky_relu = _fuse_for_leaky_relu(opr) - if leaky_relu: - matches.append(leaky_relu) - _replace_opr(net, matches) - - -def _fuse_for_conv_bias(opr): - assert ( - len(opr.out_oprs) == 1 - and isinstance(opr.out_oprs[0], Ops.ElemwiseOpr) - and opr.out_oprs[0].mode == "ADD" - ) - - bias_node = PatternNode("ADD", is_output=True) - conv_node = PatternNode(Ops.ConvolutionForwardOpr.__name__) - bias_node.inp_oprs = [conv_node] - - add_opr = opr.out_oprs[0] - if match(bias_node, add_opr): - conv_bias = Ops.ConvForwardBiasOpr( - "ConvForwardBias_" + bias_node.op.name, - conv_node.op._opr, - bias_node.inp_const[0][1], - ) - conv_bias.activation = add_opr.activation - conv_bias.inp_vars = conv_node.op.inp_vars + bias_node.op.inp_vars[1:] - conv_bias.out_vars = bias_node.op.out_vars - conv_bias.inp_oprs = conv_node.op.inp_oprs - conv_bias.out_oprs = bias_node.op.out_oprs - for node in [conv_node, bias_node]: - node.op.skip = True - return conv_bias - return None - - -@_register_tranformation_rule(TransformerRule.FUSE_FOR_CONV_BIAS) -def _(net): - """ - ConvolutionForward + Elemwise(ADD) -> ConvForwardBias - """ - matches = list() - for opr in net.all_oprs: - if ( - get_type(opr) == Ops.ConvolutionForwardOpr.__name__ - and len(opr.out_oprs) == 1 - and get_type(opr.out_oprs[0]) == "ADD" - ): - conv_bias = _fuse_for_conv_bias(opr) - if conv_bias: - matches.append(conv_bias) - _replace_opr(net, matches) - - -@_register_tranformation_rule(TransformerRule.RESHAPE_BIAS_TO_1DIM) -def _(net): - for opr in net.all_oprs: - if isinstance(opr, Ops.ConvForwardBiasOpr) and opr.inp_vars[2].ndim != 1: - bias = opr.inp_vars[2] - assert bias.shape == (1, bias.shape[1], 1, 1), ( - "bias.shape = %s" % bias.shape - ) - bias.np_data = bias.np_data.reshape(-1) - bias.shape = bias.np_data.shape - bias.ndim = 1 - - -def _fuse_for_deconv_bias(opr): - assert ( - len(opr.out_oprs) == 1 - and isinstance(opr.out_oprs[0], Ops.ElemwiseOpr) - and opr.out_oprs[0].mode == "ADD" - ) - - bias_node = PatternNode("ADD", is_output=True) - conv_node = PatternNode(Ops.ConvolutionBackwardDataOpr.__name__) - bias_node.inp_oprs = [conv_node] - - add_opr = opr.out_oprs[0] - if match(bias_node, add_opr): - deconv_bias = Ops.ConvolutionBackwardDataBiasOpr( - "ConvolutionBackwardDataBias_" + bias_node.op.name, - conv_node.op._opr, - bias_node.inp_const[0][1], - ) - deconv_bias.activation = add_opr.activation - deconv_bias.inp_vars = conv_node.op.inp_vars + bias_node.op.inp_vars[1:] - deconv_bias.out_vars = bias_node.op.out_vars - deconv_bias.inp_oprs = conv_node.op.inp_oprs - deconv_bias.out_oprs = bias_node.op.out_oprs - for node in [conv_node, bias_node]: - node.op.skip = True - return deconv_bias - - return None - - -@_register_tranformation_rule(TransformerRule.FUSE_FOR_DECONV_BIAS) -def _(net): - """ - ConvolutionBackwardData + Elemwise(ADD) -> ConvolutionBackwardDataBias - """ - matches = list() - for opr in net.all_oprs: - if ( - get_type(opr) == Ops.ConvolutionBackwardDataOpr.__name__ - and len(opr.out_oprs) == 1 - and get_type(opr.out_oprs[0]) == "ADD" - ): - deconv_bias = _fuse_for_deconv_bias(opr) - if deconv_bias: - matches.append(deconv_bias) - _replace_opr(net, matches) - - -def _fuse_for_fully_connected(opr): - assert ( - len(opr.out_oprs) == 1 - and isinstance(opr.out_oprs[0], Ops.ElemwiseOpr) - and opr.out_oprs[0].mode == "ADD" - ) - bias_node = PatternNode("ADD", is_output=True) - matrix_mul_node = PatternNode(Ops.MatrixMulOpr.__name__) - bias_node.inp_oprs = [matrix_mul_node] - - add_opr = opr.out_oprs[0] - if match(bias_node, add_opr): - fully_connected = Ops.FullyConnectedOpr( - "FullyConnected_" + bias_node.op.name, - matrix_mul_node.op._opr, - bias_node.inp_const[0][1], - ) - fully_connected.activation = add_opr.activation - fully_connected.inp_vars = ( - matrix_mul_node.op.inp_vars + bias_node.op.inp_vars[1:] - ) - fully_connected.out_vars = bias_node.op.out_vars - fully_connected.inp_oprs = matrix_mul_node.op.inp_oprs - fully_connected.out_oprs = bias_node.op.out_oprs - for node in [matrix_mul_node, bias_node]: - node.op.skip = True - return fully_connected - return None - - -@_register_tranformation_rule(TransformerRule.FUSE_FOR_FULLY_CONNECTED) -def _(net): - """ - MatrixMul + Elemwise(ADD) -> FullyConnected - """ - matches = list() - for opr in net.all_oprs: - if ( - get_type(opr) == Ops.MatrixMulOpr.__name__ - and len(opr.out_oprs) == 1 - and get_type(opr.out_oprs[0]) == "ADD" - ): - fc = _fuse_for_fully_connected(opr) - if fc: - matches.append(fc) - _replace_opr(net, matches) - - -def _replace_opr(net, matches: List[Ops.MgeOpr]): - """ - Recieve a list of :class:`~.Ops.MgeOpr`. - For each operator in :attr:`matches`, this function will insert it and its id. - - At the end, delete the orignal operators who matchs the transform rule. - """ - for opr in matches: - max_idx = max(net._opr_ids.index(i.id) for i in opr.inp_oprs) - net._opr_ids.insert(max_idx + 1, opr.id) - net.all_oprs.insert(max_idx + 1, opr) - new_idxs = [] - new_oprs = [] - for idx, opr in zip(net._opr_ids, net.all_oprs): - if opr.skip: - continue - new_idxs.append(idx) - new_oprs.append(opr) - net._opr_ids = new_idxs - net.all_oprs = new_oprs diff --git a/mgeconvert/onnx_converter/onnx_op.py b/mgeconvert/onnx_converter/onnx_op.py deleted file mode 100644 index 28c2d25..0000000 --- a/mgeconvert/onnx_converter/onnx_op.py +++ /dev/null @@ -1,752 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import numpy as np -import onnx - -from ..mge_context import ( - AxisAddRemoveOpr, - BatchNormForwardOpr, - BroadcastOpr, - ConcatOpr, - ConvolutionBackwardDataOpr, - ConvolutionBackwardFilterOpr, - ConvolutionForwardOpr, - DimshuffleOpr, - ElemwiseOpr, - GetVarShapeOpr, - Host2DeviceCopyOpr, - IdentityOpr, - MarkNoBroadcastElemwiseOpr, - MatrixMulOpr, - MultipleDeviceTensorHolderOpr, - PoolingForwardOpr, - ReduceOpr, - ReshapeOpr, - SharedDeviceTensorOpr, - SubtensorOpr, - TypeCvtOpr, - get_symvar_value, -) - -mge2onnx_dtype_mapping = { - # pylint: disable=no-member - np.float32: onnx.TensorProto.FLOAT, - np.float16: onnx.TensorProto.FLOAT16, - np.int8: onnx.TensorProto.INT8, - np.int16: onnx.TensorProto.INT16, - np.int32: onnx.TensorProto.INT32, - np.int64: onnx.TensorProto.INT64, - np.uint8: onnx.TensorProto.UINT8, -} - -MGE2ONNX = {} - - -def _register_op(*oprs): - def callback(impl): - for opr in oprs: - MGE2ONNX[opr] = impl - return impl - - return callback - - -opset_version = 8 - - -def set_opset_version(version): - global opset_version # pylint: disable=W0603 - opset_version = version - - -class OperatorBaseConverter: - - __opr_type__ = "OperatorBaseConverter" - - def __init__(self, opr): - """ - :param opr: the operator that converter converts. - :type opr: subclass of :class:`.MgeOpr` - """ - self._opr = opr - self._net_sources = [] - self._parameters = [] - - def _get_inputs(self, exclude_idx=None): - """ - Returns the names of inputs of onnx operator. - """ - if exclude_idx is None: - exclude_idx = [] - for idx, inp in enumerate(self._opr.inp_vars): - if idx not in exclude_idx: - if self._opr.inp_vars[idx].np_data is not None: - inp_tensor = onnx.helper.make_tensor_value_info( - inp.name, mge2onnx_dtype_mapping[inp.dtype], inp.shape - ) - param = onnx.numpy_helper.from_array(inp.np_data, inp.name) - self._net_sources.append(inp_tensor) - self._parameters.append(param) - - return [var.name for var in self._opr.inp_vars] - - def _get_outputs(self): - """ - Returns the names of outputs of onnx operator. - """ - return [var.name for var in self._opr.out_vars] - - def _get_attrs(self): - """ - Returns extra attributes needed by :method:`.convert` - """ - return {} - - def convert(self): - """ - Convert owning operator to onnx operator. Could be override by - subclass. - - Returns tuple (nodes, net_sources, parameters) - """ - nodes = [ - onnx.helper.make_node( - self.__opr_type__, - self._get_inputs(), - self._get_outputs(), - name=self._opr.name, - **self._get_attrs(), - ) - ] - return nodes, self._net_sources, self._parameters - - -@_register_op(MultipleDeviceTensorHolderOpr, SharedDeviceTensorOpr) -class IgnoredOperatorConverter(OperatorBaseConverter): - def convert(self): - return [], [], [] - - -@_register_op(GetVarShapeOpr) -class GetVarShapeConverter(OperatorBaseConverter): - __opr_type__ = "Shape" - - def convert(self): - shape = self._opr.out_vars[0] - shape.np_data = get_symvar_value(shape._var).astype(np.int64) - shape.dtype = np.int64 - return [], [], [] - - -@_register_op(ElemwiseOpr, IdentityOpr, MarkNoBroadcastElemwiseOpr) -class ElemwiseConverter(OperatorBaseConverter): - - support_op_map = { - "ADD": "Add", - "SUB": "Sub", - "TRUE_DIV": "Div", - "MUL": "Mul", - "RELU": "Relu", - "FUSE_ADD_RELU": "FUSE_ADD_RELU", - "FUSE_MUL_ADD3": "FUSE_MUL_ADD3", - "Identity": "Identity", - "EXP": "Exp", - "TANH": "Tanh", - "SIGMOID": "Sigmoid", - "ABS": "Abs", - "LOG": "Log", - "FLOOR": "Floor", - "CEIL": "Ceil", - "POW": "Pow", - "MAX": "Max", - } - - def __init__(self, opr): - super().__init__(opr) - assert ( - opr.mode in self.support_op_map - ), "Elemwise op doesn't support mode {}, you can implement it in ElemwiseConverter".format( - opr.mode - ) - op_type = self.support_op_map[opr.mode] - self.__opr_type__ = op_type - - def convert(self): - if self.__opr_type__ == "FUSE_ADD_RELU": - inputs = self._get_inputs() - outputs = self._get_outputs() - tmp_tensor = outputs[0] + "tmp_onnx" - add = onnx.helper.make_node("Add", inputs, [tmp_tensor]) - relu = onnx.helper.make_node("Relu", [tmp_tensor], [outputs[0]]) - return [add, relu], self._net_sources, self._parameters - if self.__opr_type__ == "FUSE_MUL_ADD3": - inputs = self._get_inputs() - outputs = self._get_outputs() - tmp_tensor = outputs[0] + "tmp_onnx" - mul = onnx.helper.make_node("Mul", [inputs[0], inputs[1]], [tmp_tensor]) - add = onnx.helper.make_node("Add", [tmp_tensor, inputs[2]], [outputs[0]]) - return [mul, add], self._net_sources, self._parameters - else: - return super().convert() - - -@_register_op(SubtensorOpr) -class SubtensorConverter(OperatorBaseConverter): - - __opr_type__ = "Slice" - - def slice_version_1(self, starts, ends, axes, _, inputs, outputs): - attr = {"axes": axes, "ends": ends, "starts": starts} - slice_op = onnx.helper.make_node("Slice", inputs, outputs, **attr) - return slice_op, [], [] - - def slice_version_10(self, starts, ends, axes, steps, inputs, outputs): - op_name = self._opr.name - inputs = inputs + [ - op_name + "_begin", - op_name + "_end", - op_name + "_axis", - op_name + "_step", - ] - begin = onnx.helper.make_tensor_value_info( - inputs[1], mge2onnx_dtype_mapping[np.int32], starts.shape - ) - end = onnx.helper.make_tensor_value_info( - inputs[2], mge2onnx_dtype_mapping[np.int32], ends.shape - ) - axis = onnx.helper.make_tensor_value_info( - inputs[3], mge2onnx_dtype_mapping[np.int32], axes.shape - ) - step = onnx.helper.make_tensor_value_info( - inputs[4], mge2onnx_dtype_mapping[np.int32], steps.shape - ) - net_sources = [begin, end, axis, step] - parameters = [ - onnx.numpy_helper.from_array(starts, inputs[1]), - onnx.numpy_helper.from_array(ends, inputs[2]), - onnx.numpy_helper.from_array(axes, inputs[3]), - onnx.numpy_helper.from_array(steps, inputs[4]), - ] - Slice = onnx.helper.make_node("Slice", inputs, outputs) - return Slice, net_sources, parameters - - def convert(self): - opr = self._opr - squeeze_axis = opr.squeeze_axis - begin_param = np.array(opr.begin_param, dtype=np.int32) - end_param = np.array(opr.end_param, dtype=np.int32) - step_param = np.array(opr.step_param, dtype=np.int32) - axis_param = np.array(opr.axis, dtype=np.int32) - inputs = [self._get_inputs(exclude_idx=list(range(1, len(opr.inp_vars))))[0]] - outputs = self._get_outputs() - slice_outputs = [ - outputs[0] if len(squeeze_axis) == 0 else outputs[0] + "tmp@onnx" - ] - slice_op = None - slice_net_sources = [] - slice_param = [] - if opset_version < 10: - slice_op, slice_net_sources, slice_param = self.slice_version_1( - begin_param, end_param, axis_param, step_param, inputs, slice_outputs - ) - else: - slice_op, slice_net_sources, slice_param = self.slice_version_10( - begin_param, end_param, axis_param, step_param, inputs, slice_outputs - ) - nodes = [] - self._parameters.extend(slice_param) - self._net_sources.extend(slice_net_sources) - nodes.append(slice_op) - if len(squeeze_axis) > 0: - Squeeze = onnx.helper.make_node( - "Squeeze", slice_outputs, outputs, axes=squeeze_axis - ) - nodes.append(Squeeze) - - return (nodes, self._net_sources, self._parameters) - - -@_register_op(DimshuffleOpr) -class DimshuffleConverter(OperatorBaseConverter): - - __opr_type__ = "Transpose" - - def _get_attrs(self): - return {"perm": list(self._opr.pattern)} - - -@_register_op(MatrixMulOpr) -class MatrixMulConvert(OperatorBaseConverter): - def convert(self): - opr = self._opr - const_0 = opr.name + "_const_0_onnx" - const_0_tensor = onnx.helper.make_tensor_value_info( - const_0, mge2onnx_dtype_mapping[np.float32], [1] - ) - const_0_param = onnx.numpy_helper.from_array( - np.array([0]).astype("float32"), const_0 - ) - inputs = self._get_inputs() - outputs = self._get_outputs() - gemm = onnx.helper.make_node( - "Gemm", - [inputs[0], inputs[1], const_0], - [outputs[0]], - alpha=1.0, - beta=0.0, - transA=opr.transposeA, - transB=opr.transposeB, - ) - - return ( - [gemm], - self._net_sources + [const_0_tensor], - self._parameters + [const_0_param], - ) - - -@_register_op(ReshapeOpr) -class ReshapeConverter(OperatorBaseConverter): - - __opr_type__ = "Reshape" - - def convert(self): - inputs = self._get_inputs(exclude_idx=[1]) - outputs = self._get_outputs() - inp_var = self._opr.inp_vars - inputs[1] = self._opr.name + "shape_onnx" - shape_tensor = onnx.helper.make_tensor_value_info( - inputs[1], mge2onnx_dtype_mapping[np.int64], inp_var[1].shape - ) - shape_param = onnx.numpy_helper.from_array( - self._opr.shape_param.astype(np.int64), inputs[1] - ) - reshape = onnx.helper.make_node("Reshape", inputs, outputs) - return ( - [reshape], - self._net_sources + [shape_tensor], - self._parameters + [shape_param], - ) - - -@_register_op(Host2DeviceCopyOpr) -class DataProviderConverter(OperatorBaseConverter): - def convert(self): - out_vars = self._opr.out_vars - inp_tensor_list = [ - onnx.helper.make_tensor_value_info( - var.name, mge2onnx_dtype_mapping[var.dtype], var.shape - ) - for var in out_vars - ] - - return [], inp_tensor_list, [] - - -@_register_op(ConvolutionForwardOpr, ConvolutionBackwardDataOpr) -class Conv2DConverter(OperatorBaseConverter): - - __opr_type__ = "Conv" - - def _get_attrs(self): - opr = self._opr - return { - "kernel_shape": [opr.kh, opr.kw], - "pads": [opr.ph, opr.pw, opr.ph, opr.pw], - "strides": [opr.sh, opr.sw], - "dilations": [opr.dilation_h, opr.dilation_w], - "group": opr.group if opr.group is not None else 1, - } - - def convert(self): - opr = self._opr - attrs = self._get_attrs() - nodes = [] - exclude_idx = [0] if attrs["group"] != 1 else [] - inputs = self._get_inputs(exclude_idx) - if isinstance(self._opr, ConvolutionBackwardDataOpr): - inputs = [inputs[1], inputs[0]] - - outputs = self._get_outputs() - if attrs["group"] != 1: - w_idx = 0 if isinstance(self._opr, ConvolutionBackwardDataOpr) else 1 - flt_shape = self._opr.inp_vars[w_idx].shape - flt_shape = [ - flt_shape[0] * flt_shape[1], - flt_shape[2], - flt_shape[3], - flt_shape[4], - ] - - if opr.param_W is not None: - inputs[1] = opr.name + "_filter_reshape_onnx" - flt = opr.param_W - flt_data = flt.reshape(flt_shape) - flt_tensor = onnx.helper.make_tensor_value_info( - inputs[1], mge2onnx_dtype_mapping[flt.dtype.type], flt_shape - ) - flt_param = onnx.numpy_helper.from_array(flt_data, inputs[1]) - self._net_sources.append(flt_tensor) - self._parameters.append(flt_param) - else: - reshape_inputs = [inputs[1], opr.name + "shape_onnx"] - shape_tensor = onnx.helper.make_tensor_value_info( - reshape_inputs[1], - mge2onnx_dtype_mapping[np.int64], - (len(flt_shape),), - ) - shape_param = onnx.numpy_helper.from_array( - np.array(flt_shape, dtype="int64"), reshape_inputs[1] - ) - self._net_sources.append(shape_tensor) - self._parameters.append(shape_param) - reshape = onnx.helper.make_node( - "Reshape", reshape_inputs, [opr.name + "_filter_reshape_onnx"] - ) - inputs[1] = opr.name + "_filter_reshape_onnx" - nodes.append(reshape) - onnx_op = "Conv" - if isinstance(self._opr, ConvolutionBackwardDataOpr): - onnx_op = "ConvTranspose" - conv2d = onnx.helper.make_node(onnx_op, inputs, [outputs[0]], **attrs) - nodes.append(conv2d) - return (nodes, self._net_sources, self._parameters) - - -@_register_op(ConvolutionBackwardFilterOpr) -class Conv2DBackwardFilterConverter(OperatorBaseConverter): - def convert(self): - opr = self._opr - # src, grad_out, weight = self._opr.inp_vars - - nodes = [] - inputs = self._get_inputs() - outputs = self._get_outputs() - - # Tile - # grad_out: (no, co, ho, wo) -> (no, co x ci / group, ho, wo) - grad_out_tile_in = outputs[0] + "_grad_out_tile_in" - grad_out_tile_source = onnx.helper.make_tensor_value_info( - grad_out_tile_in, mge2onnx_dtype_mapping[np.int64], (4,) - ) - grad_out_tile_param = onnx.numpy_helper.from_array( - np.array([1, opr.ci // opr.group, 1, 1]), grad_out_tile_in - ) - self._net_sources.append(grad_out_tile_source) - self._parameters.append(grad_out_tile_param) - grad_out_tile_out = outputs[0] + "_grad_out_tile_out" - grad_out_tile = onnx.helper.make_node( - "Tile", [inputs[1], grad_out_tile_in], [grad_out_tile_out] - ) - nodes.append(grad_out_tile) - - # Reshape - # grad_out: (no, co x ci / group, ho, wo) -> (no x co x ci / group, 1, ho, wo) - grad_out_reshape_in = outputs[0] + "_grad_out_reshape_in" - grad_out_reshape_source = onnx.helper.make_tensor_value_info( - grad_out_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) - ) - grad_out_reshape_param = onnx.numpy_helper.from_array( - np.array([opr.no * opr.co * opr.ci // opr.group, 1, opr.ho, opr.wo]), - grad_out_reshape_in, - ) - self._net_sources.append(grad_out_reshape_source) - self._parameters.append(grad_out_reshape_param) - grad_out_reshape_out = outputs[0] + "_grad_out_reshape_out" - reshape = onnx.helper.make_node( - "Reshape", [grad_out_tile_out, grad_out_reshape_in], [grad_out_reshape_out] - ) - nodes.append(reshape) - - # Reshape - # src: (ni, ci, hi, wi) -> (1, ni x ci, hi, wi) - src_reshape_in = outputs[0] + "_src_reshape_in" - src_reshape_source = onnx.helper.make_tensor_value_info( - src_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) - ) - src_reshape_param = onnx.numpy_helper.from_array( - np.array([1, opr.ni * opr.ci, opr.hi, opr.wi]), src_reshape_in - ) - self._net_sources.append(src_reshape_source) - self._parameters.append(src_reshape_param) - src_reshape_out = outputs[0] + "_src_reshape_out" - reshape = onnx.helper.make_node( - "Reshape", [inputs[0], src_reshape_in], [src_reshape_out] - ) - nodes.append(reshape) - - # Conv: - # group = ni * ci - # src(1, ni x ci, hi, wi) + grad_out(no x co x ci / group, 1, ho, wo) - # -> grad_weight(1, no x co x ci / group, ?, ?) - grad_weight = outputs[0] + "_grad_weight" - grad_weight_conv = onnx.helper.make_node( - "Conv", - [src_reshape_out, grad_out_reshape_out], - [grad_weight], - kernel_shape=[opr.ho, opr.wo], - strides=[opr.dilate_h, opr.dilate_w], - pads=[opr.ph, opr.pw, opr.ph, opr.pw], - dilations=[opr.sh, opr.sw], - group=opr.ci * opr.ni, - ) - nodes.append(grad_weight_conv) - - # Slice - # grad_weight: (1, no x co x ci // group, ?, ?) -> (1, no x co x ci // group, kh, kw) - grad_weight_slice_out = outputs[0] + "_grad_weight_slice_out" - grad_weight_slice = onnx.helper.make_node( - "Slice", - [grad_weight], - [grad_weight_slice_out], - axes=[2, 3], - starts=[0, 0], - ends=[opr.kh, opr.kw], - ) - nodes.append(grad_weight_slice) - - # Reshape - # grad_weight: (1, no x co x ci // group, kh, kw) -> (no, co x ci // group, kh, kw) - grad_weight_reshape_in = outputs[0] + "_grad_weight_reshape_in" - grad_weight_reshape_source = onnx.helper.make_tensor_value_info( - grad_weight_reshape_in, mge2onnx_dtype_mapping[np.int64], (4,) - ) - grad_weight_reshape_param = onnx.numpy_helper.from_array( - np.array([opr.no, opr.co * opr.ci // opr.group, opr.kh, opr.kw]), - grad_weight_reshape_in, - ) - self._net_sources.append(grad_weight_reshape_source) - self._parameters.append(grad_weight_reshape_param) - grad_weight_reshape_out = outputs[0] + "_grad_weight_reshape_out" - reshape = onnx.helper.make_node( - "Reshape", - [grad_weight_slice_out, grad_weight_reshape_in], - [grad_weight_reshape_out], - ) - nodes.append(reshape) - - # ReduceSum - # grad_weight: (no, co x ci // group, kh, kw) -> (1, co x ci // goup, kh, kw) - grad_weight_reduce_out = outputs[0] + "_grad_weight_reduce_out" - grad_weight_reduce = onnx.helper.make_node( - "ReduceSum", [grad_weight_reshape_out], [grad_weight_reduce_out], axes=[0], - ) - nodes.append(grad_weight_reduce) - - # Reshape - # grad_weight: (1, co x ci // group, kh, kw) -> (ci // group, co, kh, kw) - grad_weight_reshape2_in = outputs[0] + "_grad_weight_reshape2_in" - grad_weight_reshape2_source = onnx.helper.make_tensor_value_info( - grad_weight_reshape2_in, mge2onnx_dtype_mapping[np.int64], (4,) - ) - grad_weight_reshape2_param = onnx.numpy_helper.from_array( - np.array([opr.ci // opr.group, opr.co, opr.kh, opr.kw]), - grad_weight_reshape2_in, - ) - self._net_sources.append(grad_weight_reshape2_source) - self._parameters.append(grad_weight_reshape2_param) - grad_weight_reshape2_out = outputs[0] + "_grad_weight_reshape2_out" - reshape = onnx.helper.make_node( - "Reshape", - [grad_weight_reduce_out, grad_weight_reshape2_in], - [grad_weight_reshape2_out], - ) - nodes.append(reshape) - - # Transpose - grad_weight_transpose_out = outputs[0] - transpose = onnx.helper.make_node( - "Transpose", - [grad_weight_reshape2_out], - [grad_weight_transpose_out], - perm=[1, 0, 2, 3], - ) - nodes.append(transpose) - - return (nodes, self._net_sources, self._parameters) - - -@_register_op(PoolingForwardOpr) -class Pooling2DConverter(OperatorBaseConverter): - support_op_map = { - "AVERAGE": "AveragePool", - "AVERAGE_COUNT_EXCLUDE_PADDING": "AveragePool", - "MAX": "MaxPool", - } - - def __init__(self, opr): - super().__init__(opr) - assert ( - opr.mode in self.support_op_map - ), "Pooling op doesn't support mode {}, you can implement it in Pooling2DConverter".format( - opr.mode - ) - self.exclude_pad = opr.mode == "AVERAGE_COUNT_EXCLUDE_PADDING" - self.__opr_type__ = self.support_op_map[opr.mode] - - def _get_attrs(self): - opr = self._opr - attribute = { - "kernel_shape": [opr.kh, opr.kw], - "pads": [opr.ph, opr.pw, opr.ph, opr.pw], - "strides": [opr.sh, opr.sw], - } - - if self.__opr_type__ == "AveragePool": - attribute["count_include_pad"] = 0 if self.exclude_pad else 1 - - return attribute - - -@_register_op(BatchNormForwardOpr) -class BatchnormConverter(OperatorBaseConverter): - def convert(self): - inputs = self._get_inputs(exclude_idx=[1, 2, 3, 4]) - outputs = self._get_outputs() - scale_ = self._opr.scale - bias_ = self._opr.bias - mean_ = self._opr.mean - var_ = self._opr.var - inputs[1] = self._opr.name + "scale_onnx" - inputs[2] = self._opr.name + "bias_onnx" - inputs[3] = self._opr.name + "mean_onnx" - inputs[4] = self._opr.name + "var_onnx" - scale = onnx.helper.make_tensor_value_info( - inputs[1], mge2onnx_dtype_mapping[self._opr.inp_vars[1].dtype], scale_.shape - ) - bias = onnx.helper.make_tensor_value_info( - inputs[2], mge2onnx_dtype_mapping[self._opr.inp_vars[2].dtype], bias_.shape - ) - mean = onnx.helper.make_tensor_value_info( - inputs[3], mge2onnx_dtype_mapping[self._opr.inp_vars[3].dtype], mean_.shape - ) - var = onnx.helper.make_tensor_value_info( - inputs[4], mge2onnx_dtype_mapping[self._opr.inp_vars[4].dtype], var_.shape - ) - self._parameters.extend( - [ - onnx.numpy_helper.from_array(scale_, inputs[1]), - onnx.numpy_helper.from_array(bias_, inputs[2]), - onnx.numpy_helper.from_array(mean_, inputs[3]), - onnx.numpy_helper.from_array(var_, inputs[4]), - ] - ) - - bn = onnx.helper.make_node( - "BatchNormalization", inputs, [outputs[self._opr.output_idx]] - ) - return ([bn], self._net_sources + [scale, bias, mean, var], self._parameters) - - -@_register_op(ConcatOpr) -class ConcatConverter(OperatorBaseConverter): - __opr_type__ = "Concat" - - def __init__(self, opr): - super().__init__(opr) - if opset_version < 11: - assert ( - self._opr.axis >= 0 - ), "opset {} doesn't support negative aixs in concat opr".format( - opset_version - ) - - def _get_attrs(self): - return {"axis": self._opr.axis} - - -@_register_op(ReduceOpr) -class ReduceConverter(OperatorBaseConverter): - support_op_map = { - "MAX": "ReduceMax", - "SUM": "ReduceSum", - "SUM_SQR": "ReduceSumSquare", - } - - def __init__(self, opr): - super().__init__(opr) - assert ( - opr.mode in self.support_op_map - ), "Reduce op doesn't support mode {}, you can implement it in ReduceConverter".format( - opr.mode - ) - self.__opr_type__ = self.support_op_map[opr.mode] - - def _get_attrs(self): - return {"axes": [self._opr.axis]} - - def convert(self): - inputs = self._get_inputs() - outputs = self._get_outputs() - nodes = onnx.helper.make_node( - self.__opr_type__, [inputs[0]], outputs, **self._get_attrs() - ) - return [nodes], self._net_sources, self._parameters - - -@_register_op(AxisAddRemoveOpr) -class AxisAddRemoveConverter(OperatorBaseConverter): - def convert(self): - inputs = self._get_inputs() - outputs = self._get_outputs() - add_axis = [] - remove_axis = [] - for desc in self._opr.desc: - if desc["method"] == 0: - add_axis.append(desc["axisnum"]) - else: - remove_axis.append(desc["axisnum"]) - - if len(add_axis) > 0 and len(remove_axis) > 0: - assert ( - False - ), "AsixAddRemove converter doesn't support add and remove axises concurrently" - - elif len(add_axis) > 0: - unsqueeze = onnx.helper.make_node( - "Unsqueeze", inputs, outputs, axes=add_axis - ) - ret = [unsqueeze] - elif len(remove_axis) > 0: - squeeze = onnx.helper.make_node( - "Squeeze", inputs, outputs, axes=remove_axis - ) - ret = [squeeze] - else: - ret = [] - return ret, self._net_sources, self._parameters - - -@_register_op(BroadcastOpr) -class BroadcastOprConverter(OperatorBaseConverter): - def convert(self): - assert opset_version > 7, "onnx support Expand (broadcast) since opset 8" - inputs = self._get_inputs() - typecvt_node = onnx.helper.make_node( - "Cast", - [inputs[1]], - [inputs[1] + "_int64"], - to=mge2onnx_dtype_mapping[np.int64], - ) - inputs[1] = inputs[1] + "_int64" - outputs = self._get_outputs() - broadcast_node = onnx.helper.make_node("Expand", inputs, outputs) - return [typecvt_node, broadcast_node], self._net_sources, self._parameters - - -@_register_op(TypeCvtOpr) -class TypeCvtOprConverter(OperatorBaseConverter): - def convert(self): - inputs = self._get_inputs() - outputs = self._get_outputs() - target_dtype = self._opr.out_vars[0].dtype - node = onnx.helper.make_node( - "Cast", inputs, outputs, to=mge2onnx_dtype_mapping[target_dtype] - ) - return [node], self._net_sources, self._net_sources diff --git a/mgeconvert/tflite_converter/__init__.py b/mgeconvert/tflite_converter/__init__.py deleted file mode 100644 index 9cee1fb..0000000 --- a/mgeconvert/tflite_converter/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from .tflite_converter import convert_to_tflite diff --git a/mgeconvert/tflite_converter/init.sh b/mgeconvert/tflite_converter/init.sh deleted file mode 100755 index ecf22ac..0000000 --- a/mgeconvert/tflite_converter/init.sh +++ /dev/null @@ -1,41 +0,0 @@ -#!/bin/bash -e -basepath=$(cd `dirname $0`; pwd) - -sudo python3 -m pip uninstall tensorflow flatbuffers -y -sudo rm -rf /usr/local/bin/flatc -sudo rm -rf /usr/local/lib/libflatbuffers* - -# build flatbuffers -echo "building flatbuffers..." -sudo rm -rf /tmp/flatbuffers -git clone https://github.com/google/flatbuffers.git -b v1.12.0 /tmp/flatbuffers -cd /tmp/flatbuffers -cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release -DFLATBUFFERS_BUILD_SHAREDLIB=on -DCMAKE_INSTALL_PREFIX=/usr/local -make; sudo make install - -export PATH=$PATH:/usr/local/bin -# build tflite interface from schema.fbs -echo "building tflite schema..." -cd /tmp -wget https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/lite/schema/schema.fbs -sudo flatc --python schema.fbs -sudo chmod 755 /tmp/tflite -cp -r /tmp/tflite $basepath - -# build pyflatbuffers -echo "building pyflexbuffers..." -cd /tmp/flatbuffers/python -sudo python3 setup.py install - -sudo python3 -m pip install pybind11==2.6.2 - -# try to find libflatbuffers.so -export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib - -# using pyflexbuffers -cd $basepath/pyflexbuffers -PYBIND11_HEADER=$(python3 -c "import pybind11; print(pybind11.get_include())") -PYTHON_INCLUDE=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['include'])") -PYTHON_STDLIB=$(python3 -c "import sysconfig; print(sysconfig.get_paths()['stdlib'])") - -g++ fbconverter.cc --std=c++14 -fPIC --shared -I${PYBIND11_HEADER} -I${PYTHON_INCLUDE} -L${PYTHON_STDLIB} -lflatbuffers -o fbconverter.so diff --git a/mgeconvert/utils/convert_caffe.py b/mgeconvert/utils/convert_caffe.py deleted file mode 100644 index cc76c5a..0000000 --- a/mgeconvert/utils/convert_caffe.py +++ /dev/null @@ -1,47 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import argparse - -from mgeconvert.caffe_converter import convert_to_caffe - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-i", "--input", required=True, type=str, help="Input megengine dump model file" - ) - parser.add_argument( - "-c", "--prototxt", required=True, type=str, help="Output caffe .prototxt file" - ) - parser.add_argument( - "-b", - "--caffemodel", - required=True, - type=str, - help="Output caffe .caffemodel file", - ) - - parser.add_argument( - "--end_point", - default=None, - type=str, - help="end_point is used to specify which part of the mge model should be converted", - ) - - args = parser.parse_args() - outspec = None - if args.end_point is not None: - outspec = args.end_point.split(";") - convert_to_caffe( - args.input, prototxt=args.prototxt, caffemodel=args.caffemodel, outspec=outspec - ) - - -if __name__ == "__main__": - main() diff --git a/mgeconvert/utils/convert_cambricon.py b/mgeconvert/utils/convert_cambricon.py deleted file mode 100644 index 8cab2c7..0000000 --- a/mgeconvert/utils/convert_cambricon.py +++ /dev/null @@ -1,44 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import argparse - -from mgeconvert.cambricon_converter import convert_to_cambricon - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-i", "--input", required=True, type=str, help="megengine dumped model file" - ) - parser.add_argument( - "-o", "--output", required=True, type=str, help="converted Cambricon model file" - ) - parser.add_argument( - "-b", "--batch-size", default=4, type=int, help="best practice: 4" - ) - parser.add_argument("-c", "--core-number", default=1, type=int, help="c <= 16") - parser.add_argument( - "-t", "--data-type", default="float32", type=str, help="float32, float16" - ) - parser.add_argument("--use-nhwc", action="store_true", help="default nchw") - - args = parser.parse_args() - - convert_to_cambricon( - args.input, - args.output, - args.batch_size, - args.core_number, - args.data_type, - args.use_nhwc, - ) - - -if __name__ == "__main__": - main() diff --git a/mgeconvert/utils/convert_onnx.py b/mgeconvert/utils/convert_onnx.py deleted file mode 100644 index e8edcfb..0000000 --- a/mgeconvert/utils/convert_onnx.py +++ /dev/null @@ -1,45 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import argparse - -from mgeconvert.onnx_converter import convert_to_onnx - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-i", "--input", required=True, type=str, help="Input megengine dump model file" - ) - parser.add_argument( - "-o", "--output", required=True, type=str, help="Output onnx .onnx file" - ) - parser.add_argument("--opset", default=8, type=int, help="Onnx opset version") - parser.add_argument("--graph", default="graph", type=str, help="Onnx graph name") - parser.add_argument( - "--end_point", - default=None, - type=str, - help="end_point is used to specify which part of the mge model should be converted", - ) - args = parser.parse_args() - outspec = None - if args.end_point is not None: - outspec = args.end_point.split(";") - - convert_to_onnx( - args.input, - args.output, - graph_name=args.graph, - opset=args.opset, - outspec=outspec, - ) - - -if __name__ == "__main__": - main() diff --git a/mgeconvert/utils/convert_tflite.py b/mgeconvert/utils/convert_tflite.py deleted file mode 100644 index 804b66a..0000000 --- a/mgeconvert/utils/convert_tflite.py +++ /dev/null @@ -1,47 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import argparse - -from mgeconvert.tflite_converter import convert_to_tflite - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-i", "--input", required=True, type=str, help="megengine dumped model file" - ) - parser.add_argument( - "-o", "--output", required=True, type=str, help="converted TFLite model file" - ) - parser.add_argument( - "--graph-name", - default="graph0", - type=str, - help="default subgraph name in TFLite model", - ) - parser.add_argument( - "-b", "--batch-size", default=1, type=int, help="default value: 1" - ) - parser.add_argument( - "--mtk", action="store_true", help="If target flatform is MTK(P70, P80)" - ) - - args = parser.parse_args() - - convert_to_tflite( - mge_fpath=args.input, - output=args.output, - graph_name=args.graph_name, - batch_size=args.batch_size, - mtk=args.mtk, - ) - - -if __name__ == "__main__": - main() diff --git a/mgeconvert/version.py b/mgeconvert/version.py index f4a54af..8888132 100644 --- a/mgeconvert/version.py +++ b/mgeconvert/version.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,7 +5,7 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -__version__ = "0.4.2" +__version__ = "0.5.0" # required megengine version range -MEGENGINE_LOWER = "0.6.0" +MEGENGINE_LOWER = "1.0.0" diff --git a/setup.py b/setup.py index 4496a5b..e7f7bb3 100644 --- a/setup.py +++ b/setup.py @@ -13,22 +13,25 @@ targets = [] +tfversion = None class install(_install): user_options = _install.user_options + [ ("targets=", None, ""), + ("tfversion=", None, "the version of tflite schema"), ] def initialize_options(self): _install.initialize_options(self) self.targets = None + self.tfversion = None def finalize_options(self): _install.finalize_options(self) def run(self): - options = ["caffe", "onnx", "cambricon", "tflite"] + options = ["caffe", "onnx", "tflite"] if self.targets == "all": targets.extend(options) else: @@ -36,9 +39,20 @@ def run(self): with open("mgeconvert/__init__.py", "a+") as init_file: [ - init_file.write("from .%s_converter import convert_to_%s\n" % (i, i)) + init_file.write( + "from .converters.mge_to_%s import mge_to_%s\n" % (i, i) + ) for i in targets ] + [ + init_file.write( + "from .converters.tm_to_%s import tracedmodule_to_%s\n" % (i, i) + ) + for i in targets + ] + + global tfversion + tfversion = self.tfversion _install.run(self) @@ -55,7 +69,10 @@ def find_extension(self, name): raise TypeError("can not build %s" % name) def build_all(self, ext): - subprocess.check_call(ext.script) + if ext.name == "tflite" and tfversion is not None: + subprocess.check_call([ext.script, tfversion]) + else: + subprocess.check_call(ext.script) if ext.artifacts is not None: self.copy_tree(ext.artifacts, os.path.join(self.build_lib, ext.artifacts)) @@ -70,19 +87,14 @@ def __init__(self, name, script, artifacts=None): ext_modules = [ BuildExtension( name="caffe", - script="mgeconvert/caffe_converter/init.sh", - artifacts="mgeconvert/caffe_converter/caffe_pb", - ), - BuildExtension(name="onnx", script="mgeconvert/onnx_converter/init.sh"), - BuildExtension( - name="cambricon", - script="mgeconvert/cambricon_converter/init.sh", - artifacts="mgeconvert/cambricon_converter/lib/cnlib", + script="mgeconvert/backend/ir_to_caffe/init.sh", + artifacts="mgeconvert/backend/ir_to_caffe/caffe_pb", ), + BuildExtension(name="onnx", script="mgeconvert/backend/ir_to_onnx/init.sh"), BuildExtension( name="tflite", - script="mgeconvert/tflite_converter/init.sh", - artifacts="mgeconvert/tflite_converter/", + script="mgeconvert/backend/ir_to_tflite/init.sh", + artifacts="mgeconvert/backend/ir_to_tflite/", ), ] diff --git a/test/__init__.py b/test/__init__.py index 1207b5d..93515f3 100644 --- a/test/__init__.py +++ b/test/__init__.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. diff --git a/test/convolution-backward-filter.mge b/test/convolution-backward-filter.mge deleted file mode 100644 index 3fd76fa..0000000 --- a/test/convolution-backward-filter.mge +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b79baaec8b3536a3cbcab1776aa984931883c84a0a822f67b09833b60a606e6b -size 1324 diff --git a/test/mge/__init__.py b/test/mge/__init__.py new file mode 100644 index 0000000..93515f3 --- /dev/null +++ b/test/mge/__init__.py @@ -0,0 +1,7 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. diff --git a/test/mge/convolution-backward-filter.mge b/test/mge/convolution-backward-filter.mge new file mode 100644 index 0000000000000000000000000000000000000000..29fe32435fb58179ed7263f8a25427f6c72b5294 GIT binary patch literal 1324 zcmZuwYe*DP7`@wWwwjvK#)_^frnQDY3`{|ECM8Tt*}#m@LfasP$aK+=5)qRUi$b$B zF#TgtLG&=AVs}i~5_%CKVi;r)k3 zkawpwO6%U%Nb!wlm6$J1#s2(+G;lP2u=pxyE zr%IV_{h+LQzfXQ#F|17awkqp8-zd|Uax?p__ms)8Y^9?-Kwec}B#oFReAn|&XC5l= zRNQA=83i5RrGNuTl55-~pKA#kcxCk}YYT&9u2nS$qYv1{K z%-#L^aFj4X4+7$VRw4$Y()+?WH>gH6{{YzD;6G%Y=_b5*+16-A!tm|xMHqefkmBgK<1;FIUPYX$1SW5!siUdjy>1&NTpN5|1Q0*r5fMy7hlgOoYMNy6F%7N zAHh~U!X5fxEat=>*zjWvF4jXU=;aM_0G;$89lo-ud18}&&enE9sI_@J^>Pa zlFPg(4bYwgjvoexh((2*wO$#{xQtoTe=Ohxu*@m=JWCkrO>F)aLVwJck~1!a&m8oY zI>FiuV5{xB^jY$~81?XseCBl==ZkB)4clG#ods+lkz>@uz7fXVNZg%)ZbL=?>mw`N z4*VN;q4H3z0~S;sjE?euwKxyP-30&jz$xq*(*V!T!(I((!>ur}LZ2TD`V2Y$CaOZn z;U&@g_w--vp+X;=ANy~_X6~GKGN8R1NDE#g "0.6.0" and mode == "avg": - return - net = PoolOpr(mode) - mge_result = dump_mge_model(net, net.data, tmp_file) - _test(net.data, mge_result, max_error) - - -def test_subtensor(): - net = SubtensorOpr(fix_batch=True) - mge_result = dump_mge_model(net, net.data, tmp_file) - _test(net.data, mge_result, max_error) - - -def test_reshape(): - net = ReshapeOpr(fix_batch=True) - mge_result = dump_mge_model(net, net.data, tmp_file) - _test(net.data, mge_result, max_error) - - -@pytest.mark.parametrize("mode", ["bn2d"]) -def test_batchnorm(mode): - net = BnOpr(mode) - data = net.data1 if mode == "bn1d" else net.data2 - mge_result = dump_mge_model(net, data, tmp_file) - _test(data, mge_result, max_error) - - -@pytest.mark.skip(reason="not trained") -def test_convbn2d(): - net = QuantizationConvBnOpr() - mge_result = dump_mge_quantization_model(net, net.data, tmp_file) - _test(net.data, mge_result, max_error) - - -@pytest.mark.skip(reason="not trained") -def test_linear(): - net = QuantizationLinearOpr() - mge_result = dump_mge_quantization_model(net, net.data, tmp_file) - _test(net.data, mge_result, max_error) - - -@pytest.mark.skip(reason="not trained") -def test_resnet18(): - """ - visit https://github.com/MegEngine/MegEngine for details. - """ - return diff --git a/test/test_transform.py b/test/test_transform.py deleted file mode 100644 index 2490df2..0000000 --- a/test/test_transform.py +++ /dev/null @@ -1,96 +0,0 @@ -# -*- coding: utf-8 -*- -# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") -# -# Copyright (c) 2014-2021 Megvii Inc. All rights reserved. -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -import megengine as mge -from megengine.jit import trace -from mgeconvert.mge_context import ( - TopologyNetwork, - TransformerRule, - optimize_for_conversion, -) - -from .utils import ActiveOpr, ConvOpr, LinearOpr - - -def dump_mge_model(net, data, fpath="test_model.mge", optimize_for_inference=False): - if mge.__version__ <= "0.6.0": - - @trace(symbolic=True) - def inference(data, *, net): - net.eval() - output = net(data) - return output - - inference.trace(data, net=net) - mge_result = inference(data, net=net).numpy() - inference.dump( - fpath, arg_names=["data"], optimize_for_inference=optimize_for_inference, - ) - return mge_result - else: - mge_result = net(mge.tensor(data)) - net.eval() - mge_result = net(mge.tensor(data)) - - @trace(symbolic=True, capture_as_const=True) - def inference(data): - net.eval() - output = net(data) - return output - - inference(mge.tensor(data)) - inference.dump( - fpath, arg_names=["data"], optimize_for_inference=optimize_for_inference, - ) - return mge_result.numpy() - - -def test_fuse_for_leaky_relu(): - net = ActiveOpr(mode="leaky_relu") - dump_mge_model(net, net.data, fpath="test_model.mge") - net = TopologyNetwork("test_model.mge") - optimize_for_conversion(net, TransformerRule.FUSE_FOR_LEAKY_RELU) - actual = list(type(opr).__name__ for opr in net.all_oprs) - desired = ["Host2DeviceCopyOpr", "LeakyReluOpr"] - assert actual == desired - - -def test_fuse_for_conv_bias(): - net = ConvOpr(mode="normal") - dump_mge_model(net, net.data, fpath="test_model.mge") - net = TopologyNetwork("test_model.mge") - optimize_for_conversion(net, TransformerRule.FUSE_FOR_CONV_BIAS) - actual = list(type(opr).__name__ for opr in net.all_oprs) - desired = ["Host2DeviceCopyOpr", "ConvForwardBiasOpr"] - assert actual == desired - - -def test_fuse_for_deconv_bias(): - net = ConvOpr(mode="transpose") - dump_mge_model(net, net.data, "test_model.mge") - net = TopologyNetwork("test_model.mge") - optimize_for_conversion(net, TransformerRule.FUSE_FOR_DECONV_BIAS) - actual = list(type(opr).__name__ for opr in net.all_oprs) - desired = [ - "Host2DeviceCopyOpr", - "ConvolutionBackwardDataBiasOpr", - "ConvolutionBackwardDataBiasOpr", - ] - assert actual == desired - - -def test_fuse_for_fully_connected(): - net = LinearOpr() - dump_mge_model(net, net.data, "test_model.mge") - net = TopologyNetwork("test_model.mge") - optimize_for_conversion(net, TransformerRule.FUSE_ACTIVATION) - optimize_for_conversion(net, TransformerRule.FUSE_FOR_FULLY_CONNECTED) - assert net.all_oprs[-1].activation == "RELU" - actual = list(type(opr).__name__ for opr in net.all_oprs) - desired = ["Host2DeviceCopyOpr", "MatrixMulOpr", "FullyConnectedOpr"] - assert actual == desired diff --git a/test/traced_module/__init__.py b/test/traced_module/__init__.py new file mode 100644 index 0000000..93515f3 --- /dev/null +++ b/test/traced_module/__init__.py @@ -0,0 +1,7 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. diff --git a/test/traced_module/test_caffe.py b/test/traced_module/test_caffe.py new file mode 100644 index 0000000..06d77d5 --- /dev/null +++ b/test/traced_module/test_caffe.py @@ -0,0 +1,249 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from test.utils import ( + ActiveOpr, + AdaptiveAvgPool2dOpr, + BnOpr, + BroadcastOpr, + ConvOpr, + DropoutOpr, + ElemwiseOpr, + FConcatOpr, + FlattenOpr, + LinearOpr, + PoolOpr, + ReduceOpr, + ReshapeOpr, + SqueezeOpr, + SubtensorOpr, + TransposeOpr, + XORNet, + XORNet_LeakyRelu, +) + +import caffe # pylint: disable=import-error +import megengine as mge +import megengine.hub +import numpy as np +import pytest +from mgeconvert.converters.tm_to_caffe import tracedmodule_to_caffe + +from .tm_utils import get_traced_module + +max_error = 1e-6 +tmp_file = "test_module" + + +def _test_convert_result(inputs, trace_module, mge_results, max_err, input_name="x"): + + tracedmodule_to_caffe( + trace_module, prototxt=tmp_file + ".txt", caffemodel=tmp_file + ".caffemodel" + ) + caffe_net = caffe.Net(tmp_file + ".txt", tmp_file + ".caffemodel", caffe.TEST) + for i in caffe_net.blobs.keys(): + if isinstance(input_name, list): + for idx, name in enumerate(input_name): + if name in i: + caffe_net.blobs[i].data[...] = inputs[idx] + break + else: + if input_name in i: + caffe_net.blobs[i].data[...] = inputs + break + out_dict = caffe_net.forward() + + if isinstance(mge_results, dict): + assert len(list(out_dict.keys())) == len(list(mge_results.keys())) + for name in mge_results.keys(): + assert name._name in out_dict.keys() + assert out_dict[name._name].shape == mge_results[name].shape + np.testing.assert_allclose( + out_dict[name._name], mge_results[name], atol=max_err + ) + else: + caffe_results = list(out_dict.values())[0] + assert caffe_results.shape == mge_results.shape + np.testing.assert_allclose( + caffe_results, mge_results, rtol=max_err, atol=max_err + ) + + +@pytest.mark.parametrize("mode", ["normal", "group", "transpose"]) +def test_conv2d(mode): + net = ConvOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_linear(): + net = LinearOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_squeeze(): + net = SqueezeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="a") + + +@pytest.mark.parametrize("mode", ["max", "avg"]) +def test_pooling(mode): + if megengine.__version__ > "0.6.0" and mode == "avg": + return + net = PoolOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize("mode", ["bn1d", "bn2d"]) +def test_batchnorm(mode): + net = BnOpr(mode) + net.eval() + data = net.data1 if mode == "bn1d" else net.data2 + tm_module, mge_result = get_traced_module(net, mge.tensor(data)) + _test_convert_result(data, tm_module, mge_result, max_error) + + +def test_subtensor(): + net = SubtensorOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_transpose(): + net = TransposeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_concat(): + net = FConcatOpr() + data = np.random.random((1, 2, 4, 5)).astype(np.float32) + list_data = [mge.tensor(data), mge.tensor(data)] + tm_module, mge_result = get_traced_module(net, list_data) + _test_convert_result( + [data, data], tm_module, mge_result, max_error, input_name=["inps_0", "inps_1"] + ) + + +def test_reshape(): + net = ReshapeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize( + "mode", ["add", "sub", "mul", "div", "abs", "exp", "log", "max", "pow"] +) +def test_elemwise(mode): + net = ElemwiseOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="a") + + +@pytest.mark.parametrize( + "mode", ["add", "sub", "mul", "div", "abs", "exp", "log", "pow"] +) +def test_elemwise_broadcast(mode): + net = ElemwiseOpr(mode) + tm_module, mge_result = get_traced_module( + net, mge.tensor(np.array([2.0]).astype("float32")) + ) + _test_convert_result( + np.array([2.0]), tm_module, mge_result, max_error, input_name="a" + ) + + +@pytest.mark.parametrize( + "mode", ["relu", "sigmoid", "tanh", "leaky_relu", "softmax", "silu",], +) +def test_active(mode): + if megengine.__version__ < "1.5.0" and mode == "silu": + return + net = ActiveOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize("mode", ["max", "sum", "mean"]) +def test_reduce(mode): + net = ReduceOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="a") + + +def test_broadcast(): + net = BroadcastOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_flatten(): + net = FlattenOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="inps") + + +def test_dropout(): + net = DropoutOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="inps") + + +def test_adapetive_avg_pool(): + net = AdaptiveAvgPool2dOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, input_name="inps") + + +@pytest.mark.parametrize( + "model", + [ + "shufflenet_v2_x0_5", + "shufflenet_v2_x1_0", + "resnet18", + "resnet50", + "resnet101", + "resnext50_32x4d", + ], +) +def test_model(model): + data = ( + np.random.randint(0, 255, 3 * 224 * 224) + .reshape((1, 3, 224, 224)) + .astype(np.float32) + ) + if megengine.__version__ < "1.1.0": + commit_id = "dc2f2cfb228a135747d083517b98aea56e7aab92" + else: + commit_id = None + net = megengine.hub.load( + "megengine/models", model, use_cache=False, commit=commit_id, pretrained=True + ) + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(data)) + _test_convert_result(data, tm_module, mge_result, 1e-2) + + +def test_xornet(): + if megengine.__version__ < "1.1.0": + return + net = XORNet() + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_leakyrelu_model(): + if megengine.__version__ < "1.1.0": + return + net = XORNet_LeakyRelu() + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) diff --git a/test/traced_module/test_onnx.py b/test/traced_module/test_onnx.py new file mode 100644 index 0000000..023c54d --- /dev/null +++ b/test/traced_module/test_onnx.py @@ -0,0 +1,228 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +from test.utils import ( + ActiveOpr, + AdaptiveAvgPool2dOpr, + BnOpr, + BroadcastOpr, + ConvOpr, + DropoutOpr, + ElemwiseOpr, + FConcatOpr, + FlattenOpr, + LinearOpr, + PoolOpr, + ReduceOpr, + RepeatOpr, + ReshapeOpr, + SqueezeOpr, + SubtensorOpr, + TransposeOpr, + TypeCvtOpr, +) + +import megengine as mge +import megengine.hub +import numpy as np +import onnxruntime as ort +import pytest +from mgeconvert.converters.tm_to_onnx import tracedmodule_to_onnx + +from .tm_utils import get_traced_module + +max_error = 1e-6 +tmp_file = "test_model" + + +def _test_convert_result( + inputs, fpath, mge_result, max_err, min_version=7, max_version=12 +): + for version in range(min_version, max_version + 1): + tracedmodule_to_onnx( + fpath, tmp_file + ".onnx", opset=version, graph_name="graph" + ) + onnx_net = ort.InferenceSession(tmp_file + ".onnx") + if isinstance(inputs, (list, tuple)): + input_dict = {} + for i, inp in enumerate(inputs): + input_name = onnx_net.get_inputs()[i].name + X_test = inp + input_dict[input_name] = X_test + pred_onx = onnx_net.run(None, input_dict)[0] + else: + input_name = onnx_net.get_inputs()[0].name + X_test = inputs + pred_onx = onnx_net.run(None, {input_name: X_test})[0] + assert pred_onx.shape == mge_result.shape + assert pred_onx.dtype == mge_result.dtype + np.testing.assert_allclose(pred_onx, mge_result, rtol=max_err, atol=max_err) + + +@pytest.mark.parametrize("mode", ["normal", "group", "transpose"]) +def test_conv2d(mode): + net = ConvOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_linear(): + net = LinearOpr() + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize("mode", ["max", "avg"]) +def test_pool(mode): + if mode == "avg": + return + net = PoolOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize("mode", ["bn1d", "bn2d"]) +def test_batchnorm(mode): + net = BnOpr(mode) + net.eval() + data = net.data1 if mode == "bn1d" else net.data2 + tm_module, mge_result = get_traced_module(net, mge.tensor(data)) + _test_convert_result(data, tm_module, mge_result, max_error) + + +def test_subtensor(): + net = SubtensorOpr() + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_transpose(): + net = TransposeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_concat(): + net = FConcatOpr() + data = np.random.random((1, 2, 4, 5)).astype(np.float32) + list_data = [mge.tensor(data), mge.tensor(data)] + tm_module, mge_result = get_traced_module(net, list_data) + _test_convert_result([data, data], tm_module, mge_result, max_error) + + +def test_softmax(): + net = ActiveOpr(mode="softmax") + data = np.random.random((10, 8)).astype(np.float32) + tm_module, mge_result = get_traced_module(net, mge.tensor(data)) + _test_convert_result(data, tm_module, mge_result, max_error) + + +def test_squeeze(): + net = SqueezeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_reshape(): + net = ReshapeOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize( + "mode", + ["add", "sub", "mul", "div", "abs", "exp", "log", "pow", "ceil", "floor", "max",], +) +def test_elemwise(mode): + net = ElemwiseOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize("mode", ["sum", "max"]) +def test_reduce(mode): + net = ReduceOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize( + "mode", ["relu", "relu6", "tanh", "sigmoid", "hswish", "hsigmoid", "silu"] +) +def test_active(mode): + if megengine.__version__ < "1.5.0" and mode == "silu": + return + net = ActiveOpr(mode) + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_broadcast(): + net = BroadcastOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error, min_version=8) + + +def test_typecvt(): + net = TypeCvtOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_flatten(): + net = FlattenOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_dropout(): + net = DropoutOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_globalpooling(): + net = AdaptiveAvgPool2dOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +def test_repeat(): + net = RepeatOpr() + tm_module, mge_result = get_traced_module(net, mge.tensor(net.data)) + _test_convert_result(net.data, tm_module, mge_result, max_error) + + +@pytest.mark.parametrize( + "model", + [ + "shufflenet_v2_x0_5", + "shufflenet_v2_x1_0", + "resnet18", + "resnet50", + "resnet101", + "resnext50_32x4d", + ], +) +def test_model(model): + data = ( + np.random.randint(0, 255, 3 * 224 * 224) + .reshape((1, 3, 224, 224)) + .astype(np.float32) + ) + if mge.__version__ < "1.1.0": + commit_id = "dc2f2cfb228a135747d083517b98aea56e7aab92" + else: + commit_id = None + net = megengine.hub.load( + "megengine/models", model, use_cache=False, commit=commit_id, pretrained=True + ) + net.eval() + tm_module, mge_result = get_traced_module(net, mge.tensor(data)) + _test_convert_result(data, tm_module, mge_result, 1e-2) diff --git a/test/traced_module/test_qat_tflite.py b/test/traced_module/test_qat_tflite.py new file mode 100644 index 0000000..2d33be9 --- /dev/null +++ b/test/traced_module/test_qat_tflite.py @@ -0,0 +1,165 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=import-error,no-name-in-module,no-member + +from test.traced_module.test_tflite import _test_convert_result +from test.utils import ConvOpr, LinearOpr + +import megengine as mge +import megengine.module as M +import numpy as np +from megengine.core.tensor import dtype +from megengine.core.tensor.dtype import _builtin_quant_dtypes +from megengine.module.quant_dequant import QuantStub +from megengine.quantization.quantize import quantize_qat +from megengine.quantization.utils import create_qparams +from megengine.traced_module.fake_quant import FakeQuantize + +from .tm_utils import get_traced_module + +max_error = 1e-6 +tmp_file = "test_model" + + +def get_qat_net(inp_dtype, net, num_inp=1, shape=(1, 16, 32, 32)): + qat_net = quantize_qat(net) + inps = [] + for _ in range(num_inp): + data1 = mge.tensor(np.random.random(shape)) * 16 + data1 = data1.astype(inp_dtype) + inp1 = mge.tensor(dtype.convert_from_qint8(data1.numpy())) + inp1.qparams.scale = mge.tensor(dtype.get_scale(inp_dtype)) + inp1.qparams.dtype_meta = dtype._builtin_quant_dtypes["qint8"] + inps.append(inp1) + return qat_net, inps + + +def test_qat_conv_qint8(): + class QConvOpr(M.Module): + def __init__(self): + super().__init__() + self.normal_conv = M.Conv2d( + 3, 30, 3, stride=(2, 3), padding=(3, 1), dilation=(2, 2), + ) + self.normal_conv.bias = mge.Parameter( + np.random.random(self.normal_conv.bias.shape).astype(np.float32) + ) + + def forward(self, x): + x = self.normal_conv(x) + return x + + net = QConvOpr() + qat_net = quantize_qat(net) + + inp_dtype = dtype.qint8(16.0 / 128) + data = mge.tensor(np.random.random((1, 3, 224, 224))) * 16 + data = data.astype(inp_dtype) + inp = mge.tensor(dtype.convert_from_qint8(data.numpy())) + inp.qparams.scale = mge.tensor(dtype.get_scale(inp_dtype)) + inp.qparams.dtype_meta = dtype._builtin_quant_dtypes["qint8"] + + traced_module, tm_result = get_traced_module(qat_net, inp) + print(traced_module.flatten().graph) + inp = inp.astype(inp_dtype) + out_dtype = traced_module.graph.outputs[0].qparams + scale = out_dtype.scale.numpy() + _test_convert_result( + inp, traced_module, tm_result, scale=scale, require_quantize=True + ) + + +def test_deconv_qint8(): + net = ConvOpr("tflite_transpose") + qat_net = quantize_qat(net) + + inp_dtype = dtype.qint8(16.0 / 128) + data = mge.tensor(np.random.random((1, 3, 64, 64))) * 16 + data = data.astype(inp_dtype) + inp = mge.tensor(dtype.convert_from_qint8(data.numpy())) + inp.qparams.scale = mge.tensor(dtype.get_scale(inp_dtype)) + inp.qparams.dtype_meta = dtype._builtin_quant_dtypes["qint8"] + + traced_module, tm_result = get_traced_module(qat_net, inp) + print(traced_module.flatten().graph) + inp = inp.astype(inp_dtype) + out_dtype = traced_module.graph.outputs[0].qparams + scale = out_dtype.scale.numpy() + _test_convert_result( + inp, traced_module, tm_result, scale=scale, require_quantize=True + ) + + +def test_linear(): + net = LinearOpr() + inp_dtype = dtype.qint8(16.0 / 128.0) + qat_net, inps = get_qat_net(inp_dtype, net, shape=(10, 100)) + traced_module, tm_result = get_traced_module(qat_net, inps[0]) + print(traced_module.flatten().graph) + out_dtype = traced_module.graph.outputs[0].qparams + scale = out_dtype.scale.numpy() + inp = inps[0].astype(inp_dtype) + _test_convert_result( + inp, traced_module, tm_result, scale=scale, require_quantize=True + ) + + +def test_add(): + class ElemwiseOpr(M.Module): + def __init__(self,): + super().__init__() + self.data = np.ones((2, 3, 224, 224)).astype(np.float32) + self.data1 = np.random.random((1, 3, 1, 1)).astype(np.float32) + self.add1 = M.Elemwise("add") + self.add2 = M.Elemwise("add") + self.add3 = M.Elemwise("add") + + scale = mge.tensor((16.0 / 128.0)) + self.quant_stub = QuantStub() + self.quant_stub.act_fake_quant = FakeQuantize( + _builtin_quant_dtypes["qint8"] + ) + self.quant_stub.act_fake_quant.set_qparams( + create_qparams( + dtype_meta=_builtin_quant_dtypes["qint8"], + scale=scale, + zero_point=None, + ) + ) + self.quant_stub1 = QuantStub() + self.quant_stub1.act_fake_quant = FakeQuantize( + _builtin_quant_dtypes["qint8"] + ) + self.quant_stub1.act_fake_quant.set_qparams( + create_qparams( + dtype_meta=_builtin_quant_dtypes["qint8"], + scale=scale, + zero_point=None, + ) + ) + + def forward(self, a): + n = self.quant_stub(mge.tensor(np.float32(10))) + data1 = self.quant_stub1(mge.tensor(self.data1)) + x = self.add1(a, n) + y = self.add2(a, data1) + z = self.add3(x, y) + return z + + net = ElemwiseOpr() + inp_dtype = dtype.qint8(16.0 / 128.0) + qat_net, inps = get_qat_net(inp_dtype, net, shape=(1, 3, 1, 1)) + traced_module, tm_result = get_traced_module(qat_net, inps[0]) + print(traced_module.flatten().graph) + out_dtype = traced_module.graph.outputs[0].qparams + scale = out_dtype.scale.numpy() + inp = inps[0].astype(inp_dtype) + _test_convert_result( + inp, traced_module, tm_result, scale=scale, require_quantize=True + ) diff --git a/test/traced_module/test_tflite.py b/test/traced_module/test_tflite.py new file mode 100644 index 0000000..e2c4799 --- /dev/null +++ b/test/traced_module/test_tflite.py @@ -0,0 +1,240 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint:disable=import-outside-toplevel, no-name-in-module,import-error +from test.utils import ( + ConvOpr, + ElemwiseOpr, + FConcatOpr, + LinearOpr, + PoolOpr, + ReduceOpr, + ReshapeOpr, + ResizeOpr, + SoftmaxOpr, + SqueezeOpr, + SubtensorOpr, + TransposeOpr, + TypeCvtOpr, +) +from typing import Sequence + +import megengine as mge +import megengine.functional as F +import megengine.module as M +import numpy as np +import pytest +from megengine.traced_module import trace_module +from mgeconvert.converters.tm_to_tflite import tracedmodule_to_tflite +from tensorflow.lite.python import interpreter # pylint: disable=import-error + +from .tm_utils import get_traced_module + +max_error = 1e-6 +tmp_file = "test_model" + + +def _test_convert_result( + inputs, + tm, + tm_result, + max_err=max_error, + nhwc=True, + nhwc2=True, + scale=1, + zero_point=0, + require_quantize=False, +): + if not isinstance(inputs, Sequence): + inputs = [ + inputs, + ] + if not isinstance(scale, Sequence): + scale = (scale,) + if not isinstance(zero_point, Sequence): + zero_point = (zero_point,) + for i, inp in enumerate(inputs): + if nhwc and inp.ndim == 4: + inputs[i] = inp.transpose((0, 2, 3, 1)) + + tracedmodule_to_tflite( + tm, output=tmp_file + ".tflite", require_quantize=require_quantize + ) + + tfl_model = interpreter.Interpreter(model_path=tmp_file + ".tflite") + tfl_model.allocate_tensors() + + input_details = tfl_model.get_input_details() + for i, inp in enumerate(inputs): + tfl_model.set_tensor(input_details[i]["index"], inp) + tfl_model.invoke() + + pred_tfl = [] + if not isinstance(scale, Sequence): + scale = (scale,) + zero_point = (zero_point,) + for index, i in enumerate(tfl_model.get_output_details()): + out = tfl_model.tensor(i["index"])() + if nhwc2 and out.ndim == 4: + out = out.transpose((0, 3, 1, 2)) + index = len(scale) - 1 if index >= len(scale) else index + out = ((out - float(zero_point[index])) * scale[index]).astype("float32") + pred_tfl.append(out) + + if not isinstance(tm_result, Sequence): + tm_result = (tm_result,) + for i, j, s in zip(tm_result, pred_tfl, scale): + assert i.shape == j.shape + assert i.dtype == j.dtype + atol = max_err if s == 1 else s + np.testing.assert_allclose(i, j, atol=atol) + print("success!") + + +@pytest.mark.parametrize("mode", ["normal", "group", "tflite_transpose"]) +def test_conv(mode): + net = ConvOpr(mode) + data = mge.tensor(np.random.random((1, 3, 224, 224)).astype(np.float32)) + traced_module, tm_result = get_traced_module(net, data) + print(traced_module.flatten().graph) + + _test_convert_result(data, traced_module, tm_result) + + +def test_reshape(): + net = ReshapeOpr(fix_batch=True) + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result, nhwc=False) + + +@pytest.mark.parametrize("mode", ["max", "min", "sum", "mean"]) +def test_reduce(mode): + net = ReduceOpr(mode=mode) + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +@pytest.mark.parametrize( + "mode", + [ + "pow", + "exp", + "min", + "max", + "add", + "div", + "sub", + "mul", + "fuse_add_relu", + "fuse_add_sigmoid", + ], +) +def test_elemwise(mode): + net = ElemwiseOpr(mode) + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +@pytest.mark.parametrize("mode", ["max", "avg"]) +def test_pooling(mode): + if mge.__version__ > "0.6.0" and mode == "avg": + return + net = PoolOpr(mode) + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_linear(): + net = LinearOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_transopse(): + net = TransposeOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_softmax(): + net = SoftmaxOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_squeeze(): + net = SqueezeOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_slice(): + net = SubtensorOpr() + tm, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(tm.flatten().graph) + _test_convert_result(mge.tensor(net.data), tm, tm_result, nhwc=False, nhwc2=False) + + +def test_typecvt(): + net = TypeCvtOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_resize(): + net = ResizeOpr() + traced_module, tm_result = get_traced_module(net, mge.tensor(net.data)) + print(traced_module.flatten().graph) + _test_convert_result(mge.tensor(net.data), traced_module, tm_result) + + +def test_F_concat(): + net = FConcatOpr() + data1 = mge.tensor(np.random.random((1, 3, 24, 24)).astype(np.float32)) + data2 = mge.tensor(np.random.random((1, 3, 24, 24)).astype(np.float32)) + traced_module, tm_result = get_traced_module(net, [data1, data2]) + print(traced_module.flatten().graph) + _test_convert_result([data1, data2], traced_module, tm_result) + + +def test_float_func_conv(): + class FConvOpr(M.Module): + def __init__(self): + super().__init__() + self.conv = F.conv2d + + def forward( + self, + inp, + weight, + bias=None, + stride=(1, 1), + padding=(0, 0), + dilation=(1, 1), + groups=1, + ): + x = F.conv2d(inp, weight, bias, stride, padding, dilation, groups) + return x + + net = FConvOpr() + data = mge.tensor(np.random.random((1, 16, 32, 32))).astype("float32") + weight = mge.tensor(np.random.random((32, 16, 2, 2))).astype("float32") + traced_module = trace_module(net, data, weight) + tm_result = traced_module(data, weight) + _test_convert_result([data, weight], traced_module, tm_result, max_err=1e-4) diff --git a/mgeconvert/cambricon_converter/lib/__init__.py b/test/traced_module/tm_utils.py similarity index 59% rename from mgeconvert/cambricon_converter/lib/__init__.py rename to test/traced_module/tm_utils.py index 108c92c..a2fc689 100644 --- a/mgeconvert/cambricon_converter/lib/__init__.py +++ b/test/traced_module/tm_utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,9 +5,13 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -from . import model, operators, tensor -from .cnlib import cambriconLib -cambriconLib.cnmlInit(0) +# pylint: disable=import-error,no-name-in-module -cnq = cambriconLib.cnQueue() +from megengine.traced_module import trace_module + + +def get_traced_module(net, *x): + traced_module = trace_module(net, *x) + expect = traced_module(*x) + return traced_module, expect diff --git a/test/utils.py b/test/utils.py index b91ec64..b87e603 100644 --- a/test/utils.py +++ b/test/utils.py @@ -1,4 +1,3 @@ -# -*- coding: utf-8 -*- # MegEngine is Licensed under the Apache License, Version 2.0 (the "License") # # Copyright (c) 2014-2020 Megvii Inc. All rights reserved. @@ -6,6 +5,9 @@ # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + +# pylint: disable=no-member + import random import megengine as mge @@ -112,7 +114,6 @@ def __init__(self): def forward(self, x): x = self.linear(x) x = self.linear_bias(x) - x = F.relu(x) return x @@ -178,6 +179,25 @@ def forward(self, a): return F.concat([a, a], self.concat_idx) +class FConcatOpr(M.Module): + def __init__(self): + super().__init__() + self.concat_idx = random.randint(0, 3) + + def forward(self, inps): + return F.concat(inps, self.concat_idx) + + +class MConcatOpr(M.Module): + def __init__(self): + super().__init__() + self.concat_idx = random.randint(0, 3) + self.m = M.Concat() + + def forward(self, inps, axis=0): + return self.m(inps, axis) + + class SoftmaxOpr(M.Module): def __init__(self): super().__init__() @@ -249,6 +269,10 @@ def forward(self, a): x = a + mge.tensor(self.data) y = a + mge.tensor(self.data2) z = F.maximum(x, y) + elif self.mode == "min": + x = a + mge.tensor(self.data) + y = a + mge.tensor(self.data2) + z = F.minimum(x, y) elif self.mode == "pow": z = a ** 2 @@ -304,6 +328,19 @@ def forward(self, a): return F.max(a, axis=2) +class ResizeOpr(M.Module): + def __init__(self): + super().__init__() + self.data = np.random.random((1, 2, 3, 4)).astype(np.float32) + self.out_shape = [8, 8] + self.out_shape2 = [3, 4] + + def forward(self, x): + x = F.vision.interpolate(x, size=self.out_shape, mode="bilinear") + x = F.vision.interpolate(x, size=self.out_shape2, mode="bilinear") + return x + + class ActiveOpr(M.Module): str2fun = { "relu": F.relu, @@ -311,20 +348,26 @@ class ActiveOpr(M.Module): "sigmoid": F.sigmoid, "leaky_relu": F.leaky_relu, "softmax": F.softmax, - "relu6": lambda x: F.maximum(F.minimum(x, 6), 0), + "relu6": F.relu6, + "hswish": F.hswish, + "hsigmoid": F.hsigmoid, } + if mge.__version__ >= "1.5.0": + str2fun["silu"] = F.silu def __init__(self, mode, fused=False): super().__init__() self.mode = mode self.fused = fused self.data = (np.random.random((1, 2, 3, 4)).astype(np.float32) - 0.5) * 8.0 + self.sigmoid = M.Sigmoid() + self.act = ActiveOpr.str2fun[self.mode] def forward(self, x): if self.fused: - return ActiveOpr.str2fun[self.mode](x + x) + return self.act(x + x) else: - return ActiveOpr.str2fun[self.mode](x) + return self.act(x) class BroadcastOpr(M.Module): @@ -396,3 +439,42 @@ def forward(self, x): x = self.fc2(x) x = F.leaky_relu(x) return x + + +class RepeatOpr(M.Module): + def __init__(self): + super().__init__() + self.data = np.random.random((2, 3, 4)).astype("float32") + + def forward(self, x): + x = F.repeat(x, 2, axis=1) + return x + + +class FlattenOpr(M.Module): + def __init__(self): + super().__init__() + self.data = np.random.random((1, 2, 3, 4)).astype(np.float32) + + def forward(self, inps): + return F.flatten(inps) + + +class DropoutOpr(M.Module): + def __init__(self): + super().__init__() + self.data = np.random.random((1, 2, 3, 4)).astype(np.float32) + self.drop_out = M.Dropout() + + def forward(self, inps): + return self.drop_out(inps) + + +class AdaptiveAvgPool2dOpr(M.Module): + def __init__(self): + super().__init__() + self.data = np.random.random((2, 512, 64, 64)).astype(np.float32) + self.gap = M.AdaptiveAvgPool2d((2, 2)) + + def forward(self, inps): + return self.gap(inps) diff --git a/tools/change_batch.py b/tools/change_batch.py index 916b33e..ee7991f 100644 --- a/tools/change_batch.py +++ b/tools/change_batch.py @@ -1,3 +1,10 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. import argparse import megengine.core.tensor.megbrain_graph as G diff --git a/tools/convert_and_check.py b/tools/convert_and_check.py index d4bf3a1..e8e6f20 100644 --- a/tools/convert_and_check.py +++ b/tools/convert_and_check.py @@ -1,3 +1,10 @@ +# MegEngine is Licensed under the Apache License, Version 2.0 (the "License") +# +# Copyright (c) 2014-2020 Megvii Inc. All rights reserved. +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. import argparse import io import time