tensorflow_quantization.utils

class tensorflow_quantization.utils.CreateAssetsFolders(base_experiment_directory)[source]

Create empty folders to save the original and quantized TensorFlow models and their respective ONNX models for each experiment.

The following directory structure is created: base_directory -> experiment_directory (created by add_folder method) -> (fp32 [saved_model, .onnx model]), (int8 [saved_model, .onnx model]).

add_folder(folder_name: str) None[source]

Create the experiment directory (sub-folder in the base directory passed to this class).

Parameters

folder_name (str) -- name of folder

Returns

None

class tensorflow_quantization.utils.Folder(folder_name)[source]

Folder class that tracks all files for a single experiment.

class tensorflow_quantization.utils.KerasModelTraveller(print_layer_config=False)[source]

Utility class to travel Keras model and print out detailed layer information.

get_layer_information(keras_model: keras.engine.training.Model, filter_by_class=None)[source]

Print information about all layers.

Parameters
  • keras_model (tf.keras.Model) -- Keras model

  • filter_by_class (str) -- None or array of layer.__class__ to print

Returns

None

get_layer_names(keras_model: keras.engine.training.Model, filter_by_class=None)[source]

Get name of all layers in the model.

Parameters
  • keras_model (tf.keras.Model) -- Keras model

  • filter_by_class (str) -- None or array of layer.__class__ to print

Returns

None

tensorflow_quantization.utils.convert_keras_model_to_onnx(keras_model: keras.engine.training.Model, onnx_model_path: str, opset=13) None[source]

Convert in-memory Keras model into ONNX format. Works directly with CreateAssetsFolder object path.

Parameters
  • keras_model (tf.keras.Model) -- Keras model.

  • onnx_model_path (str) -- Full path to ONNX model file.

Returns

None

tensorflow_quantization.utils.convert_saved_model_to_onnx(saved_model_dir: str, onnx_model_path: str, opset=13) None[source]

Convert Keras saved model into ONNX format. Works directly with CreateAssetsFolder object path.

Parameters
  • saved_model_dir (str) -- Path to keras saved model.

  • onnx_model_path (str) -- Full path to ONNX model file.

Returns

None

tensorflow_quantization.utils.ensure_and_clean_dir(dir_path, do_clean_dir=True) None[source]

Create a directory to save test logs

Parameters
  • dir_path (str) -- directory to create / clean.

  • do_clean_dir (bool) -- boolean indicating whether to clean the directory if it already exists (remove+create).

Returns

None

tensorflow_quantization.utils.find_my_predecessors(model: keras.engine.training.Model, current_layer_name: str) List[dict][source]

Given a layer name, find all predecessors of that layer.

Parameters
  • model (tf.keras.Model) -- Keras functional model

  • current_layer_name (str) -- name of a model layer for which predecessors has to be found.

Returns

List[dict] -- List of predecessors. Each dictionary has three keys as follows,

{'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}

Raises

AssertionError -- If model is subclassed or current_layer_name is not string.

tensorflow_quantization.utils.find_my_successors(model: keras.engine.training.Model, current_layer_name: str) List[dict][source]

Given a layer name, find all successors of that layer.

Parameters
  • model (tf.keras.Model) -- Keras functional model

  • current_layer_name (str) -- name of a model layer for which successors has to be found.

Returns

List[dict] -- List of predecessors. Each dictionary has three keys as follows,

{'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}

Raises

AssertionError -- If model is subclassed or current_layer_name is not string.