#########################################################
# Copyright (C) 2020 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Joey Chou
#########################################################
from typing import Dict, List, Tuple, Optional, Callable, Union, TypeVar
from afe.backends import Backend
from afe.core.configs import AfeProcessingConfigs
from afe.core.graph_manager import transform_irmod_to_awesomenet
from afe.core.utils import wrap_parameters_to_model_configs, wrap_parameters_to_transformer_configs
from afe.load.importers.general_importer import make_model_name_from_path
from afe.ir.net import AwesomeNet
from afe._tvm._defines import TVMIRModule
from afe.ir.tensor_type import ScalarType
_A = TypeVar('_A')
[docs]
TensorShape = Tuple[int, ...]
[docs]
NodeConversionMapping = Dict[str, Callable]
def _transform_and_convert_tvm_ir_to_awesome_net(mod: TVMIRModule,
name: str,
layout: str,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False,
framework: Optional[str] = None
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Create a GraphTransformer object and transform the TVMIRModule.
:param mod: TVMIRModule.
:param layout: str. Input data layout.
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its targeted
backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: AwesomeNet.
"""
model_configs = wrap_parameters_to_model_configs(name=name, framework=framework, layout=layout,
is_quantized=is_quantized)
transformer_configs = wrap_parameters_to_transformer_configs(
indices_to_backend_dict=index_to_backend_dict
)
processing_configs = AfeProcessingConfigs(model_configs, transformer_configs)
return transform_irmod_to_awesomenet(mod, processing_configs)
[docs]
def load_tensorflow_model(pb_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
output_names: List[str],
layout: str = 'NHWC',
convert_layout: bool = False,
custom_convert_map: Optional[Dict[str, Callable]] = None,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a TensorFlow model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param pb_file_path: Path to a .pb TensorFlow model
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param output_names: List of output names of the network eg. ['output1', 'output2']
:param layout: any variation of the characters NHWC representing Batch Size, Height, Width, and Channels
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param custom_convert_map: A custom op conversion map that maps operation names to functions. Whenever an operator
with a name found in the custom_convert_map is found in TVM, the function is called with
4 arguments:
inputs = tvm relay expression inputs to operator.
attr = list of strings to operation attributes.
params = list of strings to tvm runtime arrays that are constants in the network.
mod = The tvm irmodule containing subgraphs it uses to help construct the main graph
from tensorflow.
The function then returns the tvm relay IR expression that is inserted into the model
wherever the operation occurs.
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.tensorflow, file_paths=[pb_file_path],
input_names=list(shape_dict.keys()),
input_shapes=shape_dict,
output_names=output_names,
layout=layout, custom_convert_map=custom_convert_map)
name = make_model_name_from_path(pb_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_tensorflow2_model(pb_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
output_names: List[str],
layout: str = 'NHWC',
convert_layout: bool = False,
custom_convert_map: Optional[Dict[str, Callable]] = None,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a TensorFlow model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param pb_file_path: Path to a SavedModel dir.
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param output_names: List of output names of the network eg. ['output1', 'output2']
:param layout: any variation of the characters NHWC representing Batch Size, Height, Width, and Channels
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param custom_convert_map: A custom op conversion map that maps operation names to functions. Whenever an operator
with a name found in the custom_convert_map is found in TVM, the function is called with
4 arguments:
inputs = tvm relay expression inputs to operator.
attr = list of strings to operation attributes.
params = list of strings to tvm runtime arrays that are constants in the network.
mod = The tvm irmodule containing subgraphs it uses to help construct the main graph
from tensorflow.
The function then returns the tvm relay IR expression that is inserted into the model
wherever the operation occurs.
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.tensorflow2, file_paths=[pb_file_path],
input_names=list(shape_dict.keys()),
input_shapes=shape_dict,
output_names=output_names,
layout=layout, custom_convert_map=custom_convert_map)
name = make_model_name_from_path(pb_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_pytorch_model(pt_file_path: str,
input_names: List[str],
input_shapes: List[Tuple[int, ...]],
input_dtypes: Optional[List[ScalarType]] = None,
layout: str = 'NCHW',
convert_layout: bool = False,
custom_convert_map: Optional[Dict[str, Callable]] = None,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a PyTorch model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param pt_file_path: Path to a PyTorch file (.pt) that contains the entire model
:param input_names: List of input names. eg ['input']
:param input_shapes: List of input shapes corresponding to input names. eg. [(1, 224, 224, 3)]
:param input_dtypes: List of input datatypes corresponding to input names. eg ['float32']
:param layout: Input data layout. Default is 'NCHW'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param custom_convert_map: A custom op conversion map that maps operation names to functions. Whenever an operator
with a name found in the custom_convert_map is found in TVM, the function is called with
2 arguments:
inputs = tvm relay expression inputs to operator.
input_types = list of strings indicating the input types to the operator.
The function then returns the tvm relay IR expression that is inserted into the model
wherever the operation occurs.
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.pytorch, file_paths=[pt_file_path],
input_names=input_names,
input_shapes=input_shapes,
input_types=input_dtypes,
layout=layout, custom_convert_map=custom_convert_map)
name = make_model_name_from_path(pt_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_tflite_model(tflite_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
dtype_dict: Dict[str, ScalarType],
layout: str = 'NHWC',
convert_layout: bool = False,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a TensorFlow Lite model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param tflite_file_path: Path to a .tflite file containing the TensorFlow Lite model
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param dtype_dict: Dictionary of input names to input types eg. {'input', 'float32'}
:param layout: Input data layout. Default is 'NHWC'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.tflite, file_paths=[tflite_file_path],
input_shapes=shape_dict,
input_types=dtype_dict,
layout=layout)
name = make_model_name_from_path(tflite_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_onnx_model(onnx_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
dtype_dict: Dict[str, ScalarType],
layout: str = 'NCHW',
convert_layout: bool = False,
custom_convert_map: Optional[Dict[str, Callable]] = None,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a ONNX model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param onnx_file_path: Path to a .onnx file containing the onnx model
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param dtype_dict: Dictionary of input names to input types eg. {'input', 'float32'}
:param layout: Input data layout. Default is 'NCHW'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param custom_convert_map: A custom op conversion map that maps operation names to functions. Whenever an operator
with a name found in the custom_convert_map is found in TVM, the function is called with
3 arguments
inputs = a tvm onnx_input object which contains a dictionary of tvm function inputs
attr = a dictionary of operation attributes
params = a dictionary of all the constants in the onnx network
The function then returns the tvm relay IR expression that is inserted into the model
wherever the operation occurs.
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.onnx, file_paths=[onnx_file_path],
input_names=list(shape_dict.keys()),
input_shapes=shape_dict,
input_types=dtype_dict,
layout=layout, custom_convert_map=custom_convert_map)
name = make_model_name_from_path(onnx_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_keras_model(keras_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
layout: str = 'NCHW',
convert_layout: bool = False,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a Keras model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param keras_file_path: Path to a .h5 file containing the keras model
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param layout: Input data layout. Default is 'NCHW'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.keras, file_paths=[keras_file_path],
input_shapes=shape_dict,
layout=layout)
name = make_model_name_from_path(keras_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_caffe_model(prototxt_file_path: str,
caffemodel_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
dtype_dict: Dict[str, ScalarType],
layout: str = 'NCHW',
convert_layout: bool = False,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a caffe model into an AwesomeNet format. This function is deprecated;
use import_and_transform instead.
:param prototxt_file_path: filepath to the caffe .prototxt file
:param caffemodel_file_path: filepath to the caffe .caffemodel file
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param dtype_dict: Dictionary of input names to input types eg. {'input', 'float32'}
:param layout: Input data layout. Default is 'NCHW'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.caffe, file_paths=[prototxt_file_path, caffemodel_file_path],
input_shapes=shape_dict,
input_types=dtype_dict,
layout=layout)
name = make_model_name_from_path(caffemodel_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()
[docs]
def load_caffe2_model(init_net_file_path: str,
predict_net_file_path: str,
shape_dict: Dict[str, Tuple[int, ...]],
dtype_dict: Dict[str, ScalarType],
layout: str = 'NCHW',
convert_layout: bool = False,
return_irmod: bool = False,
index_to_backend_dict: Optional[Dict[int, Backend]] = None,
is_quantized: bool = False
) -> Union[AwesomeNet, Tuple[AwesomeNet, TVMIRModule]]:
"""
Loads a caffe2 model into an AwesomeNet format.
:param init_net_file_path: filepath to the caffe2 .pb init_net file
:param predict_net_file_path: filepath to the caffe2 .pb predict_net file
:param shape_dict: Dictionary of input names to input shapes eg. {'input', (1,224,224,3)}
:param dtype_dict: Dictionary of input names to input types eg. {'input', 'float32'}
:param layout: Input data layout. Default is 'NCHW'
:param convert_layout: Set it to True to enable TVM ConvertLayout transformation
:param return_irmod: bool. Default is False. If True, the function also returns TVM Relay IR Module representation
of the Awesomenet
:param index_to_backend_dict: Optional[Dict[int, Backend]]. A dictionary of node indices to its
targeted backend.
:param is_quantized: Whether the model is pre-quantized. Default is False
:return: An AwesomeNet
"""
from afe.driver.passes import import_and_transform
from afe.load.importers.general_importer import ModelFormat, ImporterParams
# Removed features
assert convert_layout, "Cannot import without converting layout"
assert not return_irmod, "Cannot return IR module"
ip = ImporterParams(format=ModelFormat.caffe2, file_paths=[init_net_file_path, predict_net_file_path],
input_shapes=shape_dict,
input_types=dtype_dict,
layout=layout)
name = make_model_name_from_path(init_net_file_path)
compiler_pass = import_and_transform(ip, name=name, index_to_backend_dict=index_to_backend_dict,
is_quantized=is_quantized)
return compiler_pass.run()