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