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
[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
-
-
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
dictionary of metric tensors and their update operations;
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
-
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
-
static
Module contents¶
Code for creating and loading models in TensorFlow