Support bf16 onnx model (#141)

* Support bf16 onnx model

* fix the tensor init api

* fix the trainer bf16 onnx gen problem

* remove neural_engine depence

* update the code

* remove the debug

* remove all INC dependence

* update the code

* short the line to 99

Author: a32543254

nlp_toolkit/backends/neural_engine/compile/compile.py

@@ -23,15 +23,11 @@
 4. Finally, convert them to .yaml file and .bin file for model configuration and inference.
 """
 
-import datetime
-import logging as log
-import os
-import sys
-import traceback
 from collections import OrderedDict
 from .loaders.loader import Loader
 from .extractors.extractor import Extractor
 from .sub_graph.subgraph_matcher import SubGraphMatcher
+from .graph_utils import get_model_fwk_name
 
 COMPILES = OrderedDict({
     'loader': Loader,
@@ -53,5 +49,10 @@ def start_pipeline(model, config=None):
 
 
 def compile(model, config=None):
-    model = start_pipeline(model, config=config)
+    if get_model_fwk_name(model) == 'neural engine':
+        from nlp_toolkit.backends.neural_engine.compile.graph import Graph
+        model = Graph()
+        model.graph_init(model + '/conf.yaml', model + '/model.bin')
+    else:
+        model = start_pipeline(model, config=config)
     return model

nlp_toolkit/backends/neural_engine/compile/extractors/extractor.py

@@ -17,8 +17,8 @@
 
 from .tf_extractor import TensorflowExtractor
 from .onnx_extractor import ONNXExtractor
-from neural_compressor.utils import logger
-
+from .. import logger
+from ..graph_utils import get_model_fwk_name
 
 EXTRACTORS = {
     'tensorflow': TensorflowExtractor,
@@ -33,9 +33,8 @@ class Extractor(object):
     """
 
     def __call__(self, model):
-        # framework = model.framework_specific_info['framework']
-        framework = model.framework()
+        framework = model[1]
         extractor = EXTRACTORS[framework]()
-        model = extractor(model)
+        model = extractor(model[0])
         logger.info('Extract {} model done...'.format(framework))
         return model

nlp_toolkit/backends/neural_engine/compile/extractors/onnx_extractor.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 
 
-from neural_compressor.utils import logger
+from .. import logger
 from ..graph.graph import Graph
 from ..ops.op import OPERATORS
 from ..onnx_utils import graph_node_names_details
@@ -30,7 +30,7 @@ class ONNXExtractor(object):
     and output_tensors, these tensors record the source/dest op name. All of these nodes 
     (in a list) will compose a graph, which is Graph class, as the return object.
     Args:
-        model: neural_compressor TensorflowBaseModel
+        model: ONNXModel
     Return:
         Graph: Graph class, the new graph object
 
@@ -42,12 +42,12 @@ def __call__(self, model):
         logger.info('Start to extarct onnx model ops...')
         new_graph = Graph()
         new_graph.framework = 'onnxruntime'
-        for graph_input in model.graph().input:
+        for graph_input in model.graph.input:
             op_type = 'ONNXINPUT'
             new_node = OPERATORS[op_type]()
             new_node.extract('onnxruntime', graph_input, model, graph_nodes_dict)
             new_graph.insert_nodes(len(new_graph.nodes), [new_node])
-        for node in model.nodes():
+        for node in model.graph.node:
             op_type = node.op_type
             if op_type == 'Constant':
                 continue
@@ -77,7 +77,7 @@ def __call__(self, model):
                                         for tensor in graph_node.input_tensors:
                                             if origin_tensor_name + ':0' == tensor.name: 
                                                 has_tensor = False
-                                        if pre_node in model.initializer() and has_tensor:
+                                        if pre_node in model.graph.initializer and has_tensor:
                                             from onnx.numpy_helper import to_array
                                             data = to_array(pre_node)
                                             shape = list(data.shape) if data.shape != () else [1]

nlp_toolkit/backends/neural_engine/compile/extractors/tf_extractor.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 
 
-from neural_compressor.utils import logger
+from .. import logger
 from ..graph.graph import Graph
 from ..ops.op import OPERATORS
 from ..tf_utils import graph_node_names_details
@@ -29,14 +29,13 @@ class TensorflowExtractor(object):
     and output_tensors, these tensors record the source/dest op name. All of these nodes
     (in a list) will compose a graph, which is Graph class, as the return object.
     Args:
-        model: neural_compressor TensorflowBaseModel
+        model: TensorflowModel
     Return:
         Graph: Graph class, the new graph object
 
     """
     @classmethod
     def __call__(self, model):
-
         nodes = model.graph_def.node
         graph_nodes_dict = graph_node_names_details(nodes)
         logger.info('Start to extarct tensorflow model ops...')

nlp_toolkit/backends/neural_engine/compile/graph/graph.py

@@ -17,7 +17,7 @@
 
 import re
 from collections import OrderedDict
-from neural_compressor.utils import logger
+from .. import logger
 import numpy as np
 import yaml
 import os

nlp_toolkit/backends/neural_engine/compile/graph_utils.py

@@ -15,12 +15,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from neural_compressor.utils import logger
+from . import logger 
 import copy
 import re
+import os
 import numpy as np
 from collections import namedtuple, OrderedDict
 from schema import Schema, And, Or
+import importlib
 
 
 DTYPES_DICT = {"float16": "fp16",
@@ -941,3 +943,101 @@ def pattern_mapping_conf_validation(conf_dict):
     },  ignore_extra_keys=True)
 
     return dict_schema.validate(conf_dict)
+
+class LazyImport(object):
+    """Lazy import python module till use
+       Args:
+           module_name (string): The name of module imported later
+    """
+    def __init__(self, module_name):
+        self.module_name = module_name
+        self.module = None
+
+    def __getattr__(self, name):
+        try:
+            self.module = importlib.import_module(self.module_name)
+            mod = getattr(self.module, name)
+        except:
+            spec = importlib.util.find_spec(str(self.module_name + '.' + name))
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+        return mod
+
+    def __call__(self, *args, **kwargs):
+        function_name = self.module_name.split('.')[-1]
+        module_name = self.module_name.split(f'.{function_name}')[0]
+        self.module = importlib.import_module(module_name)
+        function = getattr(self.module, function_name)
+        return function(*args, **kwargs)
+
+def get_model_fwk_name(model):
+    """Detect the input model belongs to which framework
+    Args:
+        model (string): framework name that supported by Neural Engine, 
+                        if there's no available fwk info, then return 'NA'.
+    """
+    onnx = LazyImport('onnx')
+    tf = LazyImport('tensorflow')
+    def _is_onnxruntime(model):
+        try:
+            if isinstance(model, str):
+                graph = onnx.load(model)
+                assert(len(graph.graph.node) != 0)
+            else:
+                graph = model.graph
+        except:
+            pass
+        else:
+            return 'onnxruntime'
+        return 'NA'
+
+    def _is_tensorflow(model):
+        try:
+            if isinstance(model, str):
+                graph_def = tf.compat.v1.GraphDef()
+                with open(model, 'rb') as f:
+                    graph_def.ParseFromString(f.read())
+            else:
+                graph = model.graph_def
+        except:
+            pass
+        else:
+            return 'tensorflow'
+        return 'NA'
+
+    def _is_neural_engine(model):
+        if model and os.path.isdir(model):
+            file_list = os.listdir(model)
+            is_engine = True
+            if len(file_list) == 2:
+                for file_name in file_list:
+                    file_ext= os.path.splitext(file_name)
+                    front, ext = file_ext
+                    if ext == ".yaml":
+                        is_engine &= True
+                    elif ext == ".bin":
+                        is_engine &= True
+                    else:
+                        is_engine &= False
+                        logger.error("Please Input yaml and bin for neural engine.")
+                        return 'NA'
+            else:
+                return 'NA'
+            if is_engine == True:
+                return 'neural engine'
+        else:
+            return 'NA'
+
+    if isinstance(model, str):
+        absmodel = os.path.abspath(os.path.expanduser(model))
+        assert os.path.exists(absmodel) or os.path.exists(absmodel+'.pb'), \
+            'invalid input path, the file does not exist!'
+
+    checker = [_is_onnxruntime, _is_neural_engine, _is_tensorflow]
+    for handler in checker:
+        fwk_name = handler(model)
+        if fwk_name != 'NA':
+            break
+    assert fwk_name != 'NA', 'Framework is not detected correctly from model format.'
+
+    return fwk_name

nlp_toolkit/backends/neural_engine/compile/loaders/loader.py

@@ -15,17 +15,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from neural_compressor.model.model import MODELS, get_model_fwk_name, get_model_type
-from neural_compressor.utils.utility import LazyImport
+from ..graph_utils import LazyImport, get_model_fwk_name
+
 onnx = LazyImport('onnx')
- 
+tf = LazyImport('tensorflow')
+
 class Loader(object):
     def __call__(self, model):
         framework = get_model_fwk_name(model)
         if framework =='tensorflow':
-            model_type = get_model_type(model)
-            model =MODELS[framework](model_type, model)
+            if isinstance(model, str):
+                graph = tf.Graph()
+                graph_def = tf.compat.v1.GraphDef()
+                with open(model, 'rb') as f:
+                    graph_def.ParseFromString(f.read())
+                    with graph.as_default():
+                        tf.import_graph_def(graph_def, name='')
+                config = tf.compat.v1.ConfigProto()
+                model = tf.compat.v1.Session(graph=graph, config=config)
         if framework =='onnxruntime':
-            model = onnx.load(model)
-            model =MODELS[framework](model)
-        return model
+            if isinstance(model, str):
+                model = onnx.load(model)
+        return model, framework