support-pack-wheel

MANIFEST.in

@@ -0,0 +1 @@
+recursive-include mgeconvert/backend *.sh *.cc *.h *.so *.so.1 flatc

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

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
 

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

ci/requires-test.txt

@@ -1,4 +1,5 @@
 numpy>=1.17
+tqdm
 onnx==1.7.0
 onnxoptimizer
 onnx-simplifier==0.3.6

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
-

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"):

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

mgeconvert/backend/ir_to_tflite/init.sh

@@ -2,34 +2,22 @@
 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 [ ! -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

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,
 )

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, "<description for this custom option>"),
@@ -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="[email protected]",
-    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="[email protected]",
+        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",
+        ],
+    )