sparsezoo.models package

Submodules

sparsezoo.models.zoo module

Code for managing the search and creation of sparsezoo Model and Recipe objects

class sparsezoo.models.zoo.Zoo[source]

Bases: object

Provides static functions for loading and searching SparseZoo models and recipes

static download_model(domain: str, sub_domain: str, architecture: str, sub_architecture: Optional[str], framework: str, repo: str, dataset: str, training_scheme: Optional[str], sparse_name: str, sparse_category: str, sparse_target: Optional[str], release_version: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False, overwrite: bool = False) → sparsezoo.objects.model.Model[source]

Downloads a model from model repo

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • release_version – The sparsezoo release version for the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

  • overwrite – True to overwrite the file if it exists, False otherwise

Returns

The requested Model instance

static download_model_from_stub(stub: Union[str, sparsezoo.requests.base.ModelArgs], override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False, overwrite: bool = False) → sparsezoo.objects.model.Model[source]
Parameters

  • stub – the SparseZoo stub path to the model, can be a string path or ModelArgs object

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

  • overwrite – True to overwrite the file if it exists, False otherwise

Returns

The requested Model instance

static download_recipe(domain: str, sub_domain: str, architecture: str, sub_architecture: Optional[str], framework: str, repo: str, dataset: str, training_scheme: Optional[str], sparse_name: str, sparse_category: str, sparse_target: Optional[str], recipe_type: Optional[str] = None, release_version: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → sparsezoo.objects.model.Model[source]

Downloads a Recipe from the model repo

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • recipe_type – The recipe type; e.g. original, transfer_learn

  • release_version – The sparsezoo release version for the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Recipe instance

static download_recipe_base_framework_files(stub: Union[str, sparsezoo.requests.base.ModelArgs], recipe_type: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False, overwrite: bool = False, extensions: Optional[List[str]] = None) → List[str][source]
Parameters

  • stub – a string model stub that points to a SparseZoo model. recipe_type may be added as a stub parameter or path of path. i.e. “model/stub/path”, “zoo:model/stub/path”, “zoo:model/stub/path?recipe_type=transfer”, “zoo:model/stub/path/transfer”

  • recipe_type – the recipe_type to download if not original

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

  • overwrite – True to overwrite the file if it exists, False otherwise

  • extensions – List of file extensions to filter for. ex [‘.pth’, ‘.ptc’]. If None or empty list, all framework files are downloaded. Default is None

Returns

file path to the downloaded framework checkpoint files for the base weights of this recipe

static download_recipe_from_stub(stub: Union[str, sparsezoo.requests.base.ModelArgs], recipe_type: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False, overwrite: bool = False) → sparsezoo.objects.recipe.Recipe[source]

Downloads a recipe from stub. If the stub is a string, it may contain the recipe type as a stub parameter or part of the stub. i.e.

  • “model/stub/path”

  • “zoo:model/stub/path”,

  • “zoo:model/stub/path?recipe_type=original”,

  • “zoo:model/stub/path/transfer_learn”

Parameters

  • stub – the SparseZoo stub path to the recipe, can be a string path or ModelArgs object

  • recipe_type – the recipe_type to download if not original

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

  • overwrite – True to overwrite the file if it exists, False otherwise

Returns

The requested Recipe instance

static load_base_model_from_recipe(recipe: sparsezoo.objects.recipe.Recipe, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False)[source]

Loads the base model associated with a recipe

Parameters

  • recipe – the Recipe associated with the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Model instance

static load_model(domain: str, sub_domain: str, architecture: str, sub_architecture: Optional[str], framework: str, repo: str, dataset: str, training_scheme: Optional[str], sparse_name: str, sparse_category: str, sparse_target: Optional[str], release_version: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → sparsezoo.objects.model.Model[source]

Obtains a Model from the model repo

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • release_version – The sparsezoo release version for the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Model instance

static load_model_from_recipe(recipe: sparsezoo.objects.recipe.Recipe, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False)[source]

Loads the model associated with a recipe

Parameters

  • recipe – the Recipe associated with the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Model instance

static load_model_from_stub(stub: Union[str, sparsezoo.requests.base.ModelArgs], override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → sparsezoo.objects.model.Model[source]
Parameters

  • stub – the SparseZoo stub path to the model, can be a string path or ModelArgs object

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Model instance

static load_recipe(domain: str, sub_domain: str, architecture: str, sub_architecture: Optional[str], framework: str, repo: str, dataset: str, training_scheme: Optional[str], sparse_name: str, sparse_category: str, sparse_target: Optional[str], recipe_type: Optional[str] = None, release_version: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → sparsezoo.objects.recipe.Recipe[source]

Obtains a Recipe from the model repo

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • recipe_type – The recipe type; e.g. original, transfer_learn

  • release_version – The sparsezoo release version for the model

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Recipe instance

static load_recipe_from_stub(stub: Union[str, sparsezoo.requests.base.ModelArgs], recipe_type: Optional[str] = None, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → sparsezoo.objects.recipe.Recipe[source]

Loads a recipe from stub. If the stub is a string, it may contain the recipe type as a stub parameter. i.e.

  • “model/stub/path”

  • “zoo:model/stub/path”,

  • “zoo:model/stub/path?recipe_type=original”,

  • “zoo:model/stub/path/transfer_learn”

Parameters

  • stub – the SparseZoo stub path to the recipe, can be a string path or ModelArgs object

  • recipe_type – the recipe type to obtain if not original

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Recipe instance

static search_models(domain: str, sub_domain: str, architecture: Optional[str] = None, sub_architecture: Optional[str] = None, framework: Optional[str] = None, repo: Optional[str] = None, dataset: Optional[str] = None, training_scheme: Optional[str] = None, sparse_name: Optional[str] = None, sparse_category: Optional[str] = None, sparse_target: Optional[str] = None, release_version: Optional[str] = None, page: int = 1, page_length: int = 20, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → List[sparsezoo.objects.model.Model][source]

Obtains a list of Models matching the search parameters

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • release_version – The sparsezoo release version for the model

  • page – the page of values to get

  • page_length – the page length of values to get

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

The requested Model instance

static search_recipes(domain: str, sub_domain: str, architecture: Optional[str] = None, sub_architecture: Optional[str] = None, framework: Optional[str] = None, repo: Optional[str] = None, dataset: Optional[str] = None, training_scheme: Optional[str] = None, sparse_name: Optional[str] = None, sparse_category: Optional[str] = None, sparse_target: Optional[str] = None, release_version: Optional[str] = None, recipe_type: Optional[str] = None, page: int = 1, page_length: int = 20, override_folder_name: Optional[str] = None, override_parent_path: Optional[str] = None, force_token_refresh: bool = False) → List[sparsezoo.objects.recipe.Recipe][source]

Obtains a list of Recipes matching the model search parameters

Parameters

  • domain – The domain of the model the object belongs to; e.g. cv, nlp

  • sub_domain – The sub domain of the model the object belongs to; e.g. classification, segmentation

  • architecture – The architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • sub_architecture – The sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • framework – The framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • repo – The source repo for the model the object belongs to; e.g. sparseml, torchvision

  • dataset – The dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • training_scheme – The training scheme used on the model the object belongs to if any; e.g. augmented

  • sparse_name – The name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • sparse_category – The degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • sparse_target – The deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

  • release_version – The sparsezoo release version for the model

  • recipe_type – The recipe type; e.g. original, transfer_learn

  • page – the page of values to get

  • page_length – the page length of values to get

  • override_folder_name – Override for the name of the folder to save this file under

  • override_parent_path – Path to override the default save path for where to save the parent folder for this file under

  • force_token_refresh – True to refresh the auth token, False otherwise

Returns

A list of Recipe objects for models that match the given search parameters

static search_similar_models(model: Union[str, sparsezoo.requests.base.ModelArgs], match_domain: bool = True, match_sub_domain: bool = True, match_architecture: bool = True, match_sub_architecture: bool = True, match_framework: bool = True, match_repo: bool = True, match_dataset: bool = True, match_training_scheme: bool = False, match_sparse_name: bool = False, match_sparse_category: bool = False, match_sparse_target: bool = False) → List[sparsezoo.objects.model.Model][source]

Search for similar models to the given one

Parameters

  • model – The model object, a SparseZoo model stub path, or a ModelArgs object representing the base model to search similar models of

  • match_domain – True to match similar models to the current domain of the model the object belongs to; e.g. cv, nlp

  • match_sub_domain – True to match similar models to the current sub domain of the model the object belongs to; e.g. classification, segmentation

  • match_architecture – True to match similar models to the current architecture of the model the object belongs to; e.g. resnet_v1, mobilenet_v1

  • match_sub_architecture – True to match similar models to the current sub architecture (scaling factor) of the model the object belongs to; e.g. 50, 101, 152

  • match_framework – True to match similar models to the current framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • match_repo – True to match similar models to the current source repo for the model the object belongs to; e.g. sparseml, torchvision

  • match_dataset – True to match similar models to the current dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • match_training_scheme – True to match similar models to the current training scheme used on the model the object belongs to if any; e.g. augmented

  • match_sparse_name – True to match similar models to the current name describing the sparsification of the model the object belongs to, e.g. base, pruned, pruned_quant

  • match_sparse_category – True to match similar models to the current degree of sparsification of the model the object belongs to; e.g. none, conservative (~100% baseline metric), moderate (>=99% baseline metric), aggressive (<99% baseline metric)

  • match_sparse_target – True to match similar models to the current deployment target of sparsification of the model the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu

Returns

a list of models matching the given model, if any

static search_sparse_models(model: Union[str, sparsezoo.requests.base.ModelArgs], match_framework: bool = True, match_repo: bool = True, match_dataset: bool = True, match_training_scheme: bool = True) → List[sparsezoo.objects.model.Model][source]

Search for different available sparse versions based off of the current model

Parameters

  • model – The model object, a SparseZoo model stub path, or a ModelArgs object representing the base model to search different sparsifications of

  • match_framework – True to match similar models to the current framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • match_repo – True to match similar models to the current source repo for the model the object belongs to; e.g. sparseml, torchvision

  • match_dataset – True to match similar models to the current dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • match_training_scheme – True to match similar models to the current training scheme used on the model the object belongs to if any; e.g. augmented

Returns

the list of matching sparse models, if any

static search_sparse_recipes(model: Union[str, sparsezoo.requests.base.ModelArgs], recipe_type: Optional[str] = None, match_framework: bool = True, match_repo: bool = True, match_dataset: bool = True, match_training_scheme: bool = True) → List[sparsezoo.objects.recipe.Recipe][source]

Search for recipes of the given model

Parameters

  • model – The model object, a SparseZoo stub model path, or a ModelArgs object representing the base model to search for recipes

  • match_framework – True to match similar models to the current framework the model the object belongs to was trained on; e.g. pytorch, tensorflow

  • match_repo – True to match similar models to the current source repo for the model the object belongs to; e.g. sparseml, torchvision

  • match_dataset – True to match similar models to the current dataset the model the object belongs to was trained on; e.g. imagenet, cifar10

  • match_training_scheme – True to match similar models to the current training scheme used on the model the object belongs to if any; e.g. augmented

Returns

the list of matching sparsification recipes, if any

Module contents

Convenient model constructors for common models in the sparsezoo