diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..b5f6223 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include mgeconvert/backend *.sh *.cc *.h *.so *.so.1 flatc \ No newline at end of file diff --git a/ci/build_package.sh b/ci/build_package.sh new file mode 100755 index 0000000..6b49477 --- /dev/null +++ b/ci/build_package.sh @@ -0,0 +1,10 @@ +#!/bin/bash -e +basepath=$(cd `dirname $0`; pwd) + +. ${basepath}/../mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh + +cd $basepath/.. +python3 setup.py sdist +python3 setup.py bdist_wheel tflite +python3 setup.py bdist_wheel caffe +python3 setup.py bdist_wheel onnx \ No newline at end of file diff --git a/ci/pytest_tflite.sh b/ci/pytest_tflite.sh index 96699c4..5240703 100755 --- a/ci/pytest_tflite.sh +++ b/ci/pytest_tflite.sh @@ -2,10 +2,9 @@ set -e +./mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh ./mgeconvert/backend/ir_to_tflite/init.sh -# try to find libflatbuffers.so -export LD_LIBRARY_PATH=$HOME/.local/lib:$LD_LIBRARY_PATH: sudo python3 -m pip uninstall flatbuffers -y sudo python3 -m pip install tensorflow==2.5.0 diff --git a/ci/pytest_tflite_complete.sh b/ci/pytest_tflite_complete.sh index fba6ace..bdae149 100755 --- a/ci/pytest_tflite_complete.sh +++ b/ci/pytest_tflite_complete.sh @@ -2,6 +2,7 @@ set -e +./mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh ./mgeconvert/backend/ir_to_tflite/init.sh # try to find libflatbuffers.so diff --git a/ci/requires-test.txt b/ci/requires-test.txt index 80830c0..f77ed2a 100644 --- a/ci/requires-test.txt +++ b/ci/requires-test.txt @@ -1,4 +1,5 @@ numpy>=1.17 +tqdm onnx==1.7.0 onnxoptimizer onnx-simplifier==0.3.6 \ No newline at end of file diff --git a/mgeconvert/backend/ir_to_onnx/init.sh b/mgeconvert/backend/ir_to_onnx/init.sh index 605fc95..38ff8e8 100755 --- a/mgeconvert/backend/ir_to_onnx/init.sh +++ b/mgeconvert/backend/ir_to_onnx/init.sh @@ -3,4 +3,3 @@ python3 -m pip install "onnx>=1.7.0" --user python3 -m pip install onnx-simplifier --user python3 -m pip install protobuf --user - diff --git a/mgeconvert/backend/ir_to_onnx/onnx_converter.py b/mgeconvert/backend/ir_to_onnx/onnx_converter.py index 027ee5c..f2a0fa3 100644 --- a/mgeconvert/backend/ir_to_onnx/onnx_converter.py +++ b/mgeconvert/backend/ir_to_onnx/onnx_converter.py @@ -12,6 +12,7 @@ import onnx.helper import onnx.numpy_helper import onnxoptimizer # pylint: disable=import-error,no-name-in-module +from tqdm import tqdm from ...converter_ir.ir_quantizer import IRQuantizer from ...frontend.mge_to_ir.mge_utils import get_symvar_value @@ -63,7 +64,7 @@ def deduplication(inputs): _, tensor_sources, _ = _add_input_tensors(self.net.graph_inputs) inputs.extend(tensor_sources) - for opr in self.net.all_oprs: + for opr in tqdm(self.net.all_oprs): if not need_convert(opr): for tensor in opr.out_tensors: if hasattr(tensor, "_var"): diff --git a/mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh b/mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh new file mode 100755 index 0000000..bf5711b --- /dev/null +++ b/mgeconvert/backend/ir_to_tflite/build_flatbuffer.sh @@ -0,0 +1,19 @@ +#!/bin/bash -e +basepath=$(cd `dirname $0`; pwd) + + +if [ ! -d /tmp/mgeconvert ]; then + mkdir /tmp/mgeconvert +fi +TMP_DIR="/tmp/mgeconvert/flatbuffers" +rm -rf $TMP_DIR + +# build flatbuffers +git clone https://github.com/google/flatbuffers.git -b v1.12.0 $TMP_DIR +echo "building flatbuffers..." +cd $TMP_DIR +cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release -DFLATBUFFERS_BUILD_SHAREDLIB=on -DCMAKE_INSTALL_PREFIX=${basepath}/pyflexbuffers +make -j; make install +echo "build flatbuffers done" + +rm -rf $TMP_DIR diff --git a/mgeconvert/backend/ir_to_tflite/init.sh b/mgeconvert/backend/ir_to_tflite/init.sh index e0ffeef..c1448d5 100755 --- a/mgeconvert/backend/ir_to_tflite/init.sh +++ b/mgeconvert/backend/ir_to_tflite/init.sh @@ -2,8 +2,6 @@ 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 @@ -11,25 +9,15 @@ else fi echo ${FLAT_BUFFER_VERSION} -if [ ! -d /tmp/mgeconvert ]; then - mkdir /tmp/mgeconvert -fi - -TMP_DIR="/tmp/mgeconvert/flatbuffers" - -if [[ ! -L "$HOME/.local/lib/libflatbuffers.so" || "$FLATC_VERSION" != "flatc version 1.12.0" || "$FLAT_BUFFER_VERSION" != "Version: 1.12" ]]; then - rm -rf $TMP_DIR - git clone https://github.com/google/flatbuffers.git -b v1.12.0 $TMP_DIR +# install flatbuffers +if [[ "$FLAT_BUFFER_VERSION" != "Version: 1.12" ]]; then + python3 -m pip uninstall flatbuffers -y + echo "install flatbuffers..." + python3 -m pip install flatbuffers==1.12 --user fi -if [[ ! -L "$HOME/.local/lib/libflatbuffers.so" || "$FLATC_VERSION" != "flatc version 1.12.0" ]]; then - rm -rf $HOME/.local/bin/flatc - rm -rf $HOME/.local/lib/libflatbuffers* - # build flatbuffers - echo "building flatbuffers..." - cd $TMP_DIR - cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release -DFLATBUFFERS_BUILD_SHAREDLIB=on -DCMAKE_INSTALL_PREFIX=$HOME/.local - make -j; make install +if [ ! -d /tmp/mgeconvert ]; then + mkdir /tmp/mgeconvert fi @@ -44,19 +32,10 @@ if [ ! -n "$1" ] ;then fi echo "Use TFLite $tf_version!" wget https://raw.githubusercontent.com/tensorflow/tensorflow/$tf_version/tensorflow/lite/schema/schema.fbs -flatc --python schema.fbs +$basepath/pyflexbuffers/bin/flatc --python schema.fbs chmod 777 /tmp/mgeconvert/tflite cp -r /tmp/mgeconvert/tflite $basepath -# build pyflatbuffers -if [[ "$FLAT_BUFFER_VERSION" != "Version: 1.12" ]]; then - python3 -m pip uninstall flatbuffers -y - echo "building pyflexbuffers..." - export VERSION=1.12 - cd $TMP_DIR/python - python3 setup.py install --user -fi - python3 -m pip install pybind11==2.6.2 --user @@ -66,6 +45,6 @@ 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$HOME/.local/include -I${PYBIND11_HEADER} -I${PYTHON_INCLUDE} -L${PYTHON_STDLIB} -L$HOME/.local/lib -lflatbuffers -o fbconverter.so +g++ fbconverter.cc --std=c++14 -fPIC --shared -I$basepath/pyflexbuffers/include -I${PYBIND11_HEADER} -I${PYTHON_INCLUDE} -L${PYTHON_STDLIB} -L$basepath/pyflexbuffers/lib -lflatbuffers -o fbconverter.so rm -rf /tmp/mgeconvert \ No newline at end of file diff --git a/mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py b/mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py index faf9e82..68d83d6 100644 --- a/mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py +++ b/mgeconvert/backend/ir_to_tflite/pyflexbuffers/__init__.py @@ -1,6 +1,12 @@ +import ctypes import json +import os -from .fbconverter import ( # pylint: disable=import-error,no-name-in-module +libflatbuffer = os.path.join(os.path.dirname(__file__), "lib", "libflatbuffers.so.1") +ctypes.cdll.LoadLibrary(libflatbuffer) + +# pylint: disable=import-error,no-name-in-module,wrong-import-position +from .fbconverter import ( # isort:skip from_json, to_json, ) diff --git a/setup.py b/setup.py index ab70d21..5ac267c 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ import os import re import subprocess +import sys from setuptools import Extension, find_packages, setup from setuptools.command.build_ext import build_ext as _build_ext @@ -16,6 +17,25 @@ tfversion = None +def write_init(targets, tflite_schema_version=None): + with open("mgeconvert/__init__.py", "w") as init_file: + [ + 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 + ] + if "onnx" in targets: + init_file.write("from .converters.onnx_to_mge import onnx_to_mge\n") + init_file.write("from .converters.onnx_to_tm import onnx_to_tracedmodule\n") + if tflite_schema_version: + init_file.write(f"tflite_schema_version={tflite_schema_version}\n") + + class install(_install): user_options = _install.user_options + [ ("targets=", None, ""), @@ -34,27 +54,28 @@ def run(self): options = ["caffe", "onnx", "tflite"] if self.targets == "all": targets.extend(options) + elif self.targets is None: + pass else: targets.extend(i for i in options if self.targets.find(i) >= 0) - - with open("mgeconvert/__init__.py", "w") as init_file: - [ - init_file.write( - "from .converters.mge_to_%s import mge_to_%s\n" % (i, i) + write_init(targets, self.tfversion) + + if "tflite" in targets: + tflite_path = os.path.join( + os.path.dirname(__file__), "mgeconvert", "backend", "ir_to_tflite" + ) + if ( + not os.path.exists( + os.path.join(tflite_path, "pyflexbuffers", "bin", "flatc") ) - for i in targets - ] - [ - init_file.write( - "from .converters.tm_to_%s import tracedmodule_to_%s\n" % (i, i) - ) - for i in targets - ] - if "onnx" in targets: - init_file.write("from .converters.onnx_to_mge import onnx_to_mge\n") - init_file.write( - "from .converters.onnx_to_tm import onnx_to_tracedmodule\n" + or not os.path.exists( + os.path.join(tflite_path, "pyflexbuffers", "include", "flatbuffers") ) + or not os.path.exists(os.path.join(tflite_path, "pyflexbuffers", "lib")) + ): + ret = os.system(f"{tflite_path}/build_flatbuffer.sh") + if ret: + raise RuntimeError("build flatbuffer failed!") global tfversion tfversion = self.tfversion @@ -74,10 +95,11 @@ def find_extension(self, name): raise TypeError("can not build %s" % name) def build_all(self, ext): - if ext.name == "tflite" and tfversion is not None: - subprocess.check_call([ext.script, tfversion]) - else: - subprocess.check_call(ext.script) + if 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)) @@ -103,17 +125,61 @@ def __init__(self, name, script, artifacts=None): ), ] -setup( - name="mgeconvert", - version=__version__, - description="MegEngine Converter", - author="Megvii Engine Team", - author_email="brain-engine@megvii.com", - url="https://github.com/MegEngine/mgeconvert", - packages=find_packages(exclude=["test", "test.*"]), - ext_modules=ext_modules, - cmdclass={"install": install, "build_ext": build_ext}, - include_package_data=True, - install_requires=["numpy"], - scripts=["bin/convert"], -) +if __name__ == "__main__": + install_requires = ["numpy", "tqdm"] + requires_mapping = { + "onnx": ["onnx>=1.7.0", "onnx-simplifier", "protobuf",], + "caffe": ["protobuf>=3.11.1"], + "tflite": ["flatbuffers==1.12.0", "pybind11==2.6.2"], + "all": [ + "onnx>=1.7.0", + "onnx-simplifier", + "protobuf>=3.11.1", + "flatbuffers==1.12.0", + ], + } + pkg_name = "mgeconvert" + if len(sys.argv) >= 2: + if sys.argv[1] == "bdist_wheel": + if len(sys.argv) == 2: + targets.extend(["onnx", "caffe", "tflite"]) + else: + assert sys.argv[2] in ["onnx", "caffe", "tflite"] + targets.append(sys.argv[2]) + pkg_name += "-" + sys.argv[2] + sys.argv = sys.argv[:2] + for t in targets: + if t in requires_mapping: + install_requires.extend(requires_mapping[t]) + elif sys.argv[1] == "install": + assert ( + len(sys.argv) > 2 + ), 'use "--targets=[eg, tflite,caffe,onnx,all]" to indicate converters to install' + for v in sys.argv[2:]: + if v.startswith("--targets="): + target_cvts = v.split("=")[1].split(",") + for opt in target_cvts: + assert opt in requires_mapping + install_requires.extend(requires_mapping[opt]) + break + write_init(targets) + + setup( + name=pkg_name, + version=__version__, + description="MegEngine Converter", + author="Megvii Engine Team", + author_email="brain-engine@megvii.com", + url="https://github.com/MegEngine/mgeconvert", + packages=find_packages(exclude=["test", "test.*"]), + ext_modules=ext_modules, + cmdclass={"install": install, "build_ext": build_ext}, + include_package_data=True, + install_requires=install_requires, + scripts=["bin/convert"], + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + ], + )