sparseml.tensorflow_v1.models package

Subpackages

• sparseml.tensorflow_v1.models.classification package
  • Submodules
  • sparseml.tensorflow_v1.models.classification.mnist module
  • sparseml.tensorflow_v1.models.classification.mobilenet module
  • sparseml.tensorflow_v1.models.classification.mobilenet_v2 module
  • sparseml.tensorflow_v1.models.classification.resnet module
  • sparseml.tensorflow_v1.models.classification.vgg module
  • Module contents

Submodules

sparseml.tensorflow_v1.models.estimator module

Model function creator classes to be used with estimator

class sparseml.tensorflow_v1.models.estimator.ClassificationEstimatorModelFn

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

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

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])

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'>)

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]

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

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])

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

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

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

Bases: abc.ABC

Base class for model function creator

create(model_const: Callable, *args, **kwargs)

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

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

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])

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'>)

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]

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

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

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

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

Bases: object

Registry class for creating models

static available_keys() → List[str]

Returns

the keys (models) currently available in the registry

static create(key: str, *args, **kwargs) → Any

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

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

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)

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)

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')

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

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