sparseml.tensorflow_v1.models package

Submodules

sparseml.tensorflow_v1.models.estimator module

Model function creator classes to be used with estimator

class sparseml.tensorflow_v1.models.estimator. ClassificationEstimatorModelFn [source]

Bases: sparseml.tensorflow_v1.models.estimator.EstimatorModelFn

Model function creator for classification models

create_loss ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , labels : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) tensorflow.python.framework.ops.Tensor [source]

Create loss function

Parameters
  • net_outputs – output tensors of the model graph

  • labels – ground truth labels

  • params – the model function params

Returns

a loss tensor

create_metric_update_ops_hook ( metrics_dict : Dict [ str , Tuple [ tensorflow.python.framework.ops.Tensor , tensorflow.python.framework.ops.Operation ] ] , metrics_initializers_dict : Dict [ str , List [ tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) sparseml.tensorflow_v1.models.estimator.MetricUpdateOpsHook [source]

Create hooks for the update operations of the collected metrics

Parameters
  • metrics_dict – dictionary of metrics, created as a result of create_metrics function

  • params – the model function params

Returns

a SessionRunHook instance

create_metrics ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , labels : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) [source]

Create metrics for evaluation

Parameters
  • net_outputs – output tensors of the model graph

  • labels – ground truth labels

  • params – the model function params

Returns

dictionary of metrics and their reset operations

create_modifier_ops_and_update_hook ( params: Dict[str , Any]) -> (<class 'sparseml.tensorflow_v1.optim.manager.ScheduledModifierManager'> , <class 'sparseml.tensorflow_v1.optim.modifier.ModifierSessionRunHook'> ) [source]

Create modifier ops and their update hook to run

Parameters

params – the model function params

Returns

a SessionRunHook instance

create_predictions ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) Dict [ str , Any ] [source]

Create predictions used to define the estimator spec in prediction mode

Parameters
  • net_outputs – output tensors of the model graph

  • params – the model function params

Returns

dictionary of metric tensors

create_scaffold ( modifier_manager : sparseml.tensorflow_v1.optim.manager.ScheduledModifierManager , params : Dict [ str , Any ] ) tensorflow.python.training.monitored_session.Scaffold [source]

Create scaffold to be attached to the train estimator spec, containing at least the saver

Parameters

params – the model function params

Returns

a Scaffold instance

create_summary_op ( metrics_dict : Dict [ str , Tuple [ tensorflow.python.framework.ops.Tensor , tensorflow.python.framework.ops.Operation ] ] , params : Dict [ str , Any ] ) [source]

Create summary op given metric dictionary

create_train_summary_hook ( metrics_dict : Dict [ str , Tuple [ tensorflow.python.framework.ops.Tensor , tensorflow.python.framework.ops.Operation ] ] , params : Dict [ str , Any ] ) tensorflow.python.training.basic_session_run_hooks.SummarySaverHook [source]

Create hook for the summary of metrics

Parameters
  • metrics_dict – dictionary of metrics, created as a result of create_metrics function

  • params – the model function params

Returns

a SsummarySaverHook instance

create_training_op ( loss : tensorflow.python.framework.ops.Tensor , params : Dict [ str , Any ] ) tensorflow.python.framework.ops.Operation [source]

Create training op for optimization

Parameters
  • loss – the loss tensor

  • params – the model function params

Returns

an Operation minimizing loss

class sparseml.tensorflow_v1.models.estimator. EstimatorModelFn [source]

Bases: abc.ABC

Base class for model function creator

create ( model_const : Callable , * args , ** kwargs ) [source]

Create a model function to be used to create estimator

Parameters
  • model_const – function to create model graph

  • arg – additional positional arguments passed into model_const

  • kwargs – additional keyword arguments passed into model_const

abstract create_loss ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , labels : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) tensorflow.python.framework.ops.Tensor [source]

Create loss function

Parameters
  • net_outputs – output tensors of the model graph

  • labels – ground truth labels

  • params – the model function params

Returns

a loss tensor

abstract create_metric_update_ops_hook ( metrics_dict : Dict [ str , Tuple [ tensorflow.python.framework.ops.Tensor , tensorflow.python.framework.ops.Operation ] ] , metrics_initializers_dict : Dict [ str , List [ tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) sparseml.tensorflow_v1.models.estimator.MetricUpdateOpsHook [source]

Create hooks for the update operations of the collected metrics

Parameters
  • metrics_dict – dictionary of metrics, created as a result of create_metrics function

  • params – the model function params

Returns

a SessionRunHook instance

abstract create_metrics ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , labels : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) [source]

Create metrics for evaluation

Parameters
  • net_outputs – output tensors of the model graph

  • labels – ground truth labels

  • params – the model function params

Returns

  1. dictionary of metric tensors and their update operations;

  2. list of extra/internal vars created for the metrics if any

abstract create_modifier_ops_and_update_hook ( params: Dict[str , Any]) -> (<class 'sparseml.tensorflow_v1.optim.manager.ScheduledModifierManager'> , <class 'sparseml.tensorflow_v1.optim.modifier.ModifierSessionRunHook'> ) [source]

Create modifier ops and their update hook to run

Parameters

params – the model function params

Returns

a SessionRunHook instance

abstract create_predictions ( net_outputs : Union [ tensorflow.python.framework.ops.Tensor , Dict [ str , tensorflow.python.framework.ops.Tensor ] ] , params : Dict [ str , Any ] ) Dict [ str , Any ] [source]

Create predictions used to define the estimator spec in prediction mode

Parameters
  • net_outputs – output tensors of the model graph

  • params – the model function params. If “apply_softmax” is specified in params then softmax is apply to the net outputs

Returns

dictionary of metric tensors

abstract create_scaffold ( modifier_manager : sparseml.tensorflow_v1.optim.manager.ScheduledModifierManager , params : Dict [ str , Any ] ) tensorflow.python.training.monitored_session.Scaffold [source]

Create scaffold to be attached to the train estimator spec, containing at least the saver

Parameters

params – the model function params

Returns

a Scaffold instance

abstract create_train_summary_hook ( metrics_dict : Dict [ str , Tuple [ tensorflow.python.framework.ops.Tensor , tensorflow.python.framework.ops.Operation ] ] , params : Dict [ str , Any ] ) tensorflow.python.training.basic_session_run_hooks.SummarySaverHook [source]

Create hook for the summary of metrics

Parameters
  • metrics_dict – dictionary of metrics, created as a result of create_metrics function

  • params – the model function params

Returns

a SummarySaverHook instance

abstract create_training_op ( loss : tensorflow.python.framework.ops.Tensor , params : Dict [ str , Any ] ) tensorflow.python.framework.ops.Operation [source]

Create training op for optimization

Parameters
  • loss – the loss tensor

  • params – the model function params

Returns

an Operation minimizing loss

sparseml.tensorflow_v1.models.registry module

Code related to the PyTorch model registry for easily creating models.

class sparseml.tensorflow_v1.models.registry. ModelRegistry [source]

Bases: object

Registry class for creating models

static available_keys ( ) List [ str ] [source]
Returns

the keys (models) currently available in the registry

static create ( key : str , * args , ** kwargs ) Any [source]

Create a new model for the given key

Parameters
  • key – the model key (name) to create

  • args – any args to supply to the graph constructor

  • kwargs – any keyword args to supply to the graph constructor

Returns

the outputs from the created graph

static create_estimator ( key : str , model_dir : str , model_fn_params : Optional [ Dict [ str , Any ] ] , run_config : tensorflow_estimator.python.estimator.run_config.RunConfig , * args , ** kwargs ) tensorflow_estimator.python.estimator.estimator.Estimator [source]

Create Estimator for a model given the key and extra parameters

Parameters
  • key – the key that the model was registered with

  • model_dir – directory to save results

  • model_fn_params – parameters for model function

  • run_config – RunConfig used by the estimator during training

  • args – additional positional arguments to pass into model constructor

  • kwargs – additional keyword arguments to pass into model constructor

Returns

an Estimator instance

static create_zoo_model ( key : str , pretrained : Union [ bool , str ] = True , pretrained_dataset : Optional [ str ] = None ) sparsezoo.objects.model.Model [source]

Create a sparsezoo Model for the desired model in the zoo

Parameters
  • key – the model key (name) to retrieve

  • pretrained – True to load pretrained weights; to load a specific version give a string with the name of the version (pruned-moderate, base), default True

  • pretrained_dataset – The dataset to load for the model

Returns

the sparsezoo Model reference for the given model

static input_shape ( key : str ) [source]
Parameters

key – the model key (name) to create

Returns

the specified input shape for the model

static load_pretrained ( key : str , pretrained : Union [ bool , str ] = True , pretrained_dataset : Optional [ str ] = None , pretrained_path : Optional [ str ] = None , remove_dynamic_tl_vars : bool = False , sess : Optional [ tensorflow.python.client.session.Session ] = None , saver : Optional [ tensorflow.python.training.saver.Saver ] = None ) [source]

Load pre-trained variables for a given model into a session. Uses a Saver object from TensorFlow to restore the variables from an index and data file.

Parameters
  • key – the model key (name) to create

  • pretrained – True to load the default pretrained variables, a string to load a specific pretrained graph (ex: base, optim, optim-perf), or False to not load any pretrained weights

  • pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

  • pretrained_path – A path to the pretrained variables to load, if provided will override the pretrained param

  • remove_dynamic_tl_vars – True to remove the vars that are used for transfer learning (have a different shape and should not be restored), False to keep all vars in the Saver. Only used if saver is None

  • sess – The session to load the model variables into if pretrained_path or pretrained is supplied. If not supplied and required, then will use the default session

  • saver – The Saver instance to use to restore the variables for the graph if pretrained_path or pretrained is supplied. If not supplied and required, then will create one using the ModelRegistry.saver function

static register ( key : Union [ str , List [ str ] ] , input_shape : Any , domain : str , sub_domain : str , architecture : str , sub_architecture : str , default_dataset : str , default_desc : str , default_model_fn_creator : sparseml.tensorflow_v1.models.estimator.EstimatorModelFn , base_name_scope : str , tl_ignore_tens : List [ str ] , repo_source : str = 'sparseml' ) [source]

Register a model with the registry. Should be used as a decorator

Parameters
  • key – the model key (name) to create

  • input_shape – the specified input shape for the model

  • domain – the domain the model belongs to; ex: cv, nlp, etc

  • sub_domain – the sub domain the model belongs to; ex: classification, detection, etc

  • architecture – the architecture the model belongs to; ex: resnet, mobilenet, etc

  • sub_architecture – the sub architecture the model belongs to; ex: 50, 101, etc

  • default_dataset – the dataset to use by default for loading pretrained if not supplied

  • default_desc – the description to use by default for loading pretrained if not supplied

  • default_model_fn_creator – default model creator to use when creating estimator instance

  • base_name_scope – the base string used to create the graph under

  • tl_ignore_tens – a list of tensors to ignore restoring for if transfer learning

  • repo_source – the source repo for the model, default is sparseml

Returns

the decorator

static saver ( key : str , remove_dynamic_tl_vars : bool = False ) tensorflow.python.training.saver.Saver [source]

Get a tf compat saver that contains only the variables for the desired architecture specified by key. Note, the architecture must have been created in the current graph already to work.

Parameters
  • key – the model key (name) to get a saver instance for

  • remove_dynamic_tl_vars – True to remove the vars that are used for transfer learning (have a different shape and should not be restored), False to keep all vars in the Saver

Returns

a Saver object with the appropriate vars for the model to restore

Module contents

Code for creating and loading models in TensorFlow