afe.load.importers.general_importer

Attributes

`TensorShape`
`InputMapping`
`NodeConversionMapping`

Classes

`ModelFormat`	Generic enumeration.
`ImporterParams`

Functions

`make_model_name_from_path`(→ str)	Create a model name from the name of the file the model was loaded from.
`default_layout`(→ str)	Get the default layout for a model format. This is the data layout that
`detect_format`(→ ModelFormat)	Determine the format of a model stored in files based on the filenames.
`update_with_detected_format`(→ ImporterParams)	Ensure that the importer parameters specify a model format. If a format is
`import_from_import_params`(...)	Import a module into TVM.
`import_model_to_tvm`([format, file_paths, input_names, ...])	General-purpose import function that wraps the other importers and is permissive in what inputs it takes.
`onnx_source`(→ ImporterParams)	Construct ImporterParams with required values for importing an ONNX model.
`pytorch_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a PyTorch model.
`tensorflow_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a Tensorflow model.
`tensorflow2_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a Tensorflow model.
`tflite_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a TFLite model.
`keras_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a Keras model.
`caffe_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a Caffe model.
`caffe2_source`(→ ImporterParams)	Construct ImporterParams with required values for importing a Caffe2 model.

Module Contents

afe.load.importers.general_importer.TensorShape[source]

afe.load.importers.general_importer.InputMapping[source]

afe.load.importers.general_importer.NodeConversionMapping[source]

class afe.load.importers.general_importer.ModelFormat[source]

Generic enumeration.

Derive from this class to define new enumerations.

tensorflow = 'tensorflow'[source]

tensorflow2 = 'tensorflow2'[source]

tflite = 'tflite'[source]

keras = 'keras'[source]

pytorch = 'pytorch'[source]

onnx = 'onnx'[source]

caffe2 = 'caffe2'[source]

caffe = 'caffe'[source]

afe.load.importers.general_importer.make_model_name_from_path(path: str) → str[source]: Create a model name from the name of the file the model was loaded from. This is used in functions that do not take the model name as user input.

afe.load.importers.general_importer.default_layout(model_format: ModelFormat, ndim: int = 4) → str[source]: Get the default layout for a model format. This is the data layout that the format uses by convention. Returns either 4D or 5D layout.

class afe.load.importers.general_importer.ImporterParams[source]

format: ModelFormat | None[source]

file_paths: List[str][source]

input_names: List[str] | None = None[source]

output_names: List[str] | None = None[source]

input_types: None | InputMapping[afe.ir.tensor_type.ScalarType] | List[afe.ir.tensor_type.ScalarType] = None[source]

input_shapes: None | InputMapping[TensorShape] | List[TensorShape] = None[source]

layout: str | None = None[source]

custom_convert_map: NodeConversionMapping | None = None[source]

property shape_dict: Dict[afe.ir.defines.NodeName, afe.ir.defines.InputShape][source]

property dtype_dict: Dict[afe.ir.defines.NodeName, afe.ir.tensor_type.ScalarType][source]

afe.load.importers.general_importer.detect_format(file_paths: List[str]) → ModelFormat[source]

Determine the format of a model stored in files based on the filenames. Raise an exception if the format can’t be determined. This function does not try to access the files.

Parameters:: file_paths – Paths of the model files. Different formats use different numbers of files.
Returns:: The inferred format.

afe.load.importers.general_importer.update_with_detected_format(import_params: ImporterParams) → ImporterParams[source]

Ensure that the importer parameters specify a model format. If a format is not given, detect and record the format. Raises an exception if a format is not given and cannot be detected.

Parameters:: import_params – Importer parameters that may have None in its format field
Returns:: Parameters derived from import_params, having a non-None format field

afe.load.importers.general_importer.import_from_import_params(import_params: ImporterParams) → Tuple[afe._tvm._defines.TVMIRModule, List[str] | None][source]

Import a module into TVM.

Parameters:: import_params – Parameters specifying what to import
Returns:: The imported module and its outputs’ names. The output names are only returned if the source model format has output names.

afe.load.importers.general_importer.import_model_to_tvm(format: ModelFormat | None = None, file_paths: List[str] = [], input_names: List[str] | None = None, output_names: List[str] | None = None, input_types: None | InputMapping[afe.ir.tensor_type.ScalarType] | List[afe.ir.tensor_type.ScalarType] = None, input_shapes: None | InputMapping[TensorShape] | List[TensorShape] = None, layout: str | None = None, custom_convert_map: NodeConversionMapping | None = None)[source]: General-purpose import function that wraps the other importers and is permissive in what inputs it takes.

afe.load.importers.general_importer.onnx_source(model_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], dtype_dict: Dict[str, afe.ir.tensor_type.ScalarType]) → ImporterParams[source]

Construct ImporterParams with required values for importing an ONNX model.

Parameters:

model_path – Path to the model file.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
dtype_dict – Dictionary with names and scalar numeric types of the model’s input tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.pytorch_source(model_path: str, input_names: List[str], input_shapes: List[Tuple[int, Ellipsis]]) → ImporterParams[source]

Construct ImporterParams with required values for importing a PyTorch model.

Parameters:

model_path – Path to the model file.
input_names – Names of the model’s input tensors.
input_shapes – Shapes of the model’s input tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.tensorflow_source(model_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], output_names: List[str]) → ImporterParams[source]

Construct ImporterParams with required values for importing a Tensorflow model.

Parameters:

model_path – Path to the model file.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
output_names – Names of the model’s output tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.tensorflow2_source(model_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], output_names: List[str]) → ImporterParams[source]

Construct ImporterParams with required values for importing a Tensorflow model.

Parameters:

model_path – Path to the model file.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
output_names – Names of the model’s output tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.tflite_source(model_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], dtype_dict: Dict[str, afe.ir.tensor_type.ScalarType]) → ImporterParams[source]

Construct ImporterParams with required values for importing a TFLite model.

Parameters:

model_path – Path to the model file.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
dtype_dict – Dictionary with names and scalar numeric types of the model’s input tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.keras_source(model_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], layout: str) → ImporterParams[source]

Construct ImporterParams with required values for importing a Keras model.

Parameters:

model_path – Path to the model file.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
layout – Keras can be in both channel-first and channel-last layout.

Returns:

ImporterParams.

afe.load.importers.general_importer.caffe_source(prototxt_file_path: str, caffemodel_file_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], dtype_dict: Dict[str, afe.ir.tensor_type.ScalarType]) → ImporterParams[source]

Construct ImporterParams with required values for importing a Caffe model.

Parameters:

prototxt_file_path – Prototxt file path.
caffemodel_file_path – Caffe model file path.
shape_dict – Dictionary with names and shapes of the model’s input tensors.
dtype_dict – Dictionary with names and scalar numeric types of the model’s input tensors.

Returns:

ImporterParams.

afe.load.importers.general_importer.caffe2_source(init_net_file_path: str, predict_net_file_path: str, shape_dict: Dict[str, Tuple[int, Ellipsis]], dtype_dict: Dict[str, afe.ir.tensor_type.ScalarType]) → ImporterParams[source]

Construct ImporterParams with required values for importing a Caffe2 model.

Parameters:

init_net_file_path
predict_net_file_path
shape_dict – Dictionary with names and shapes of the model’s input tensors.
dtype_dict – Dictionary with names and scalar numeric types of the model’s input tensors.

Returns:

ImporterParams.