.. _Model Loading:

Loading Floating-Point Models
#############################

This section describes how to load/import supported floating-point models in the ModelSDK. The floating-point model is imported into ModelSDK using the ``load_model`` API. The ``load_model`` API requires parameters to be provided and these parameters are specific to the type of model being loaded.

ONNX Format Models
------------------

For ONNX models, we need to know the input names, their shapes and their data types. The input names are associated with their shapes and data types using dictionaries.

The input shape dictionary is made up of key, value pairs in which the key is an input name (string) and the value is a tuple of integers defining the shape.
The input types dictionary is made up of key, value pairs in which the key is an input name (string) and the value is a ScalarType (generally ScalarType.float32).

The layout is assumed to be ``NCHW``.


The example below shows how to load a floating-point ONNX model that has an input named ``input_1`` with a shape of ``1,3,224,224`` in ``NCHW`` format:

.. code-block:: python

    # imports
    from afe.load.importers.general_importer import ImporterParams, onnx_source
    from afe.apis.loaded_net import load_model
    from afe.ir.tensor_type import ScalarType


    # model path
    onnx_model_path='../path/to/onnx/model/file'

    # input shapes dictionary: each key,value pair defines an input and its shape 
    input_shapes = {'input_1': (1,3,224,224)}

    # input types dictionary: each key,value pair defines an input and its type 
    input_types = {'input_1': ScalarType.float32}

    # importer parameters
    importer_params: ImporterParams = onnx_source(model_path=onnx_model_path,
                                                shape_dict=input_shapes,
                                                dtype_dict=input_types)

    # load ONNX floating-point model into LoadedNet format
    loaded_net = load_model(importer_params)


PyTorch Format Models
---------------------

For native PyTorch models, we must specify the input names and their shapes. The input names are specified as a list of strings and the input shapes are defined as a list of tuples of integers. The association between input names and input shapes is based on the list indices. Layout defaults to ``NCHW``.

This example shows how to load a floating-point PyTorch model that has an input named  ``1`` with a shape of ``1,3,224,224`` in ``NCHW`` format:

.. code-block:: python

    # imports
    from afe.load.importers.general_importer import ImporterParams, pytorch_source
    from afe.apis.loaded_net import load_model


    # model path
    pyt_model_path='../path/to/pytorch/model/file'

    # model input names - list of strings
    input_names = ['1']

    # input shapes - list of tuples
    # assumed to be NCHW format - batchsize is required
    input_shape = [(1,3,224,224)]

    # importer parameters - layout assumed to be NCHW
    importer_params: ImporterParams = pytorch_source(model_path=pyt_model_path,
                                                    input_names=input_names,
                                                    input_shapes=input_shapes)

    # load PyTorch floating-point model into LoadedNet format
    loaded_net = load_model(importer_params)

Keras Format Models
-------------------

For the Keras models in HDF5 format, we must specify a dictionary that associates the model inputs with their shapes. The ``layout`` parameter defaults to ``NHWC`` but can be changed to ``NCHW`` using the layout parameter.

This example shows how to load a floating-point Keras model that has an input named ``input_1`` with a shape of ``1,224,224,3`` in ``NHWC`` format:

.. code-block:: python

    # imports
    from afe.load.importers.general_importer import ImporterParams, keras_source
    from afe.apis.loaded_net import load_model


    # model path
    keras_model_path='../path/to/keras/model/file'

    # input shapes dictionary: each key,value pair defines an input and its shape 
    input_shapes_dict = {'input_1': (1,224,224,3)}

    # importer parameters - layout specified as NHWC
    importer_params: ImporterParams = keras_source(model_path=keras_model_path,                                                       
                                                shape_dict=input_shapes_dict,
                                                layout='NHWC')

    # load Keras floating-point model into LoadedNet format
    loaded_net = load_model(importer_params)

Tensorflow1 Format Models
-------------------------

For the Tensorflow1 models in protobuf format, we must specify a dictionary that associates the model inputs with their shapes and provide a list of the output names. Layout defaults to ``NHWC``.

The example below shows how to load a floating-point Tensorflow1 model that has an input named ``input_1`` with a shape of ``1,224,224,3`` in ``NHWC`` format and an output named prediction:

.. code-block:: python

    # imports
    from afe.load.importers.general_importer import ImporterParams, tensorflow_source
    from afe.apis.loaded_net import load_model


    # model path
    tf_model_path='../path/to/tensorflow/model/file'

    # input shapes dictionary: each key,value pair defines an input and its shape 
    input_shapes = {'input_1': (1,224,224,3)}

    # output names - list of strings
    output_names = ['prediction']

    # importer parameters
    importer_params: ImporterParams = tensorflow_source(model_path=tf_model_path,
                                                        input_shapes=input_shapes,
                                                        output_names=output_names)

    # load Tensorflow floating-point model into LoadedNet format
    loaded_net = load_model(importer_params)


TFLite Format Models
--------------------

For the TFLite models, the input names are associated with their shapes and data types using dictionaries.

The input shape dictionary is made up of key, value pairs in which the key is an input name (string) and the value is a tuple of integers defining the shape.
The input types dictionary is made up of key, value pairs in which the key is an input name (string) and the value is a ``ScalarType`` (generally ``ScalarType.float32``).

Layout defaults to ``NHWC``.

The example below shows how to load a floating-point TFLite model that has an input named ``input_1`` with a shape of ``1,224,224,3`` in ``NHWC`` format:

.. code-block:: python
        
    # imports
    from afe.load.importers.general_importer import ImporterParams, tflite_source
    from afe.apis.loaded_net import load_model
    from afe.ir.tensor_type import ScalarType


    # model path
    tflite_model_path='../path/to/tflite/model/file'

    # input shapes dictionary: each key,value pair defines an input and its shape 
    input_shapes = {'input_1': (1,224,224,3)}

    # input types dictionary: each key,value pair defines an input and its type 
    input_types = {'input_1': ScalarType.float32}

    # importer parameters
    importer_params: ImporterParams = tflite_source(model_path=tflite_model_path,
                                                    shape_dict=input_shapes,
                                                    dtype_dict=input_types)

    # load TFLite floating-point model into LoadedNet format
    loaded_net = load_model(importer_params)

Pre-Quantized TFLite Models
---------------------------

The TFLite models that have been quantized using the TFLite quantization can also be imported into ModelSDK using exactly the same code that was used for importing the floating-point TFLite model but with one exception - the input_types dictionary will need to indicate the correct data type(int8, uint8):

.. code-block:: python

    # input types dictionary: each key-value pair defines an input and its type 
    input_types = {'input_1': ScalarType.uint8}