sparseml.keras.models package

Subpackages

• sparseml.keras.models.classification package
  • Submodules
  • sparseml.keras.models.classification.resnet module
  • Module contents
• sparseml.keras.models.external package
  • Submodules
  • sparseml.keras.models.external.keras_applications module
  • Module contents

Submodules

sparseml.keras.models.registry module

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

class sparseml.keras.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, pretrained: Union[bool, str] = False, pretrained_path: Optional[str] = None, pretrained_dataset: Optional[str] = None, **kwargs) → tensorflow.python.keras.engine.training.Model

Create a new model for the given key

Parameters

• key – the model key (name) to create

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

• pretrained_path – A model file path to load into the created model

• pretrained_dataset – The dataset to load for the model

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

Returns

the instantiated model

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 (optim, optim-perf), 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) → Any

Parameters

key – the model key (name) to create

Returns

the specified input shape for the model

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, 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

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

Returns

the decorator

static register_wrapped_model_constructor(wrapped_constructor: Callable, key: Union[str, List[str]], input_shape: Any, domain: str, sub_domain: str, architecture: str, sub_architecture: str, default_dataset: str, default_desc: str, repo_source: str)

Register a model with the registry from a model constructor or provider function

Parameters

• wrapped_constructor – Model constructor wrapped to be compatible by call from ModelRegistry.create should have pretrained, pretrained_path, pretrained_dataset, load_strict, ignore_error_tensors, and kwargs as arguments

• 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

• repo_source – the source repo for the model; ex: sparseml, torchvision

Returns

The constructor wrapper registered with the registry

Module contents

Code for creating and loading datasets in Keras