neural_compressor.tensorflow.utils.model_wrappers

The concrete model wrappers for parsing all TF model formats.

Classes

BaseModel	Base class of all neural_compressor.model, will play graph role.
TensorflowBaseModel	Build Tensorflow Base Model.
TensorflowSavedModelModel	Build Tensorflow saved model.
TensorflowLLMModel	The class Tensorflow saved model whose GraphDef exceeding maximum protobuf size of 2GB.
TensorflowCheckpointModel	Build Tensorflow checkpoint model.
KerasModel	Build Keras model.
TensorflowModel	A wrapper to construct a Tensorflow Model.

Functions

get_tf_model_type(model)	The interface of getting type of tensorflow models.
get_model_type(model)	Get Tensorflow mode type.
validate_graph_node(graph_def, node_names)	Validate nodes exist in the graph_def.
validate_and_inference_input_output(graph_def, ...)	Validate and inference the input and output tensor names of graph_def.
graph_session(model, input_tensor_names, ...)	Helper to build session with tf.compat.v1.Graph.
graph_def_session(model, input_tensor_names, ...)	Build session with tf.compat.v1.GraphDef.
frozen_pb_session(model, input_tensor_names, ...)	Build session with frozen pb.
load_saved_model(model, saved_model_tags, ...)	Load graph_def from saved model with the default serving signature key.
try_loading_keras(model, input_tensor_names, ...)	Try different ways of loading keras models.
keras_session(model, input_tensor_names, ...)	Build session with keras model.
slim_session(model, input_tensor_names, ...)	Build session with slim model.
checkpoint_session(model, input_tensor_names, ...)	Build session with ckpt model.
estimator_session(model, input_tensor_names, ...)	Build session with estimator model.
saved_model_session(model, input_tensor_names, ...)	Build session with saved model.

Module Contents

neural_compressor.tensorflow.utils.model_wrappers.get_tf_model_type(model)

The interface of getting type of tensorflow models.

neural_compressor.tensorflow.utils.model_wrappers.get_model_type(model)

Get Tensorflow mode type.

Parameters:

model (string or model object) – model path or model object.

Returns:

model type

Return type:

string

neural_compressor.tensorflow.utils.model_wrappers.validate_graph_node(graph_def, node_names)

Validate nodes exist in the graph_def.

Parameters:

• graph_def (tf.compat.v1.GraphDef) – tf.compat.v1.GraphDef object.

• node_names (list of string) – node names to be validated.

neural_compressor.tensorflow.utils.model_wrappers.validate_and_inference_input_output(graph_def, input_tensor_names, output_tensor_names)

Validate and inference the input and output tensor names of graph_def.

Parameters:

• graph_def (tf.compat.v1.GraphDef) – tf.compat.v1.GraphDef object.

• input_tensor_names (list of string) – input_tensor_names of graph_def.

• output_tensor_names (list of string) – output_tensor_names of graph_def.

Returns:

validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

input_tensor_names (list of string)

neural_compressor.tensorflow.utils.model_wrappers.graph_session(model, input_tensor_names, output_tensor_names, **kwargs)

Helper to build session with tf.compat.v1.Graph.

Parameters:

• model (tf.compat.v1.Graph) – tf.compat.v1.Graph object.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.graph_def_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with tf.compat.v1.GraphDef.

Parameters:

• model (tf.compat.v1.GraphDef) – tf.compat.v1.GraphDef object.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object input_tensor_names (list of string): validated input_tensor_names output_tensor_names (list of string): validated output_tensor_names

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.frozen_pb_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with frozen pb.

Parameters:

• model (string) – model path.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.load_saved_model(model, saved_model_tags, input_tensor_names, output_tensor_names)

Load graph_def from saved model with the default serving signature key.

Parameters:

• model – Directory of the SavedModel.

• saved_model_tags – Set of tags identifying the MetaGraphDef within the SavedModel to analyze.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

The loaded GraphDef. input_tensors: List of input tensors. output_tensors: List of output tensors.

Return type:

graph_def

neural_compressor.tensorflow.utils.model_wrappers.try_loading_keras(model, input_tensor_names, output_tensor_names)

Try different ways of loading keras models.

Parameters:

• model (string or tf.keras.Model) – model path or tf.keras.Model object.

• input_tensor_names (list of string) – input tensor names of the model.

• output_tensor_names (list of string) – output tensor names of the model.

Returns:

tf.compat.v1.Session object. input_names (list of string): validated input names. output_names (list of string): validated output names.

Return type:

graph_def (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.keras_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with keras model.

Parameters:

• model (string or tf.keras.Model) – model path or tf.keras.Model object.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.slim_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with slim model.

Parameters:

• model (string) – model path.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.checkpoint_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with ckpt model.

Parameters:

• model (string) – model path.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – validated output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.estimator_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with estimator model.

Parameters:

• model (tf.estimator.Estimator) – tf.estimator.Estimator object.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

• kwargs (dict) – other required parameters, like input_fn.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

neural_compressor.tensorflow.utils.model_wrappers.saved_model_session(model, input_tensor_names, output_tensor_names, **kwargs)

Build session with saved model.

Parameters:

• model (string) – model path.

• input_tensor_names (list of string) – input_tensor_names of model.

• output_tensor_names (list of string) – output_tensor_names of model.

Returns:

tf.compat.v1.Session object. input_tensor_names (list of string): validated input_tensor_names. output_tensor_names (list of string): validated output_tensor_names.

Return type:

sess (tf.compat.v1.Session)

class neural_compressor.tensorflow.utils.model_wrappers.BaseModel(model, **kwargs)

Base class of all neural_compressor.model, will play graph role.

class neural_compressor.tensorflow.utils.model_wrappers.TensorflowBaseModel(model, **kwargs)

Build Tensorflow Base Model.

class neural_compressor.tensorflow.utils.model_wrappers.TensorflowSavedModelModel(model, **kwargs)

Build Tensorflow saved model.

class neural_compressor.tensorflow.utils.model_wrappers.TensorflowLLMModel(model, **kwargs)

The class Tensorflow saved model whose GraphDef exceeding maximum protobuf size of 2GB.

class neural_compressor.tensorflow.utils.model_wrappers.TensorflowCheckpointModel(model, **kwargs)

Build Tensorflow checkpoint model.

class neural_compressor.tensorflow.utils.model_wrappers.KerasModel(model, **kwargs)

Build Keras model.

class neural_compressor.tensorflow.utils.model_wrappers.TensorflowModel

A wrapper to construct a Tensorflow Model.