Quick Tour

Each model in the SparseZoo has a specific stub that identifies it. The stubs are made up of the following structure:

DOMAIN/SUB_DOMAIN/ARCHITECTURE{-SUB_ARCHITECTURE}/FRAMEWORK/REPO/DATASET{-TRAINING_SCHEME}/SPARSE_NAME-SPARSE_CATEGORY-{SPARSE_TARGET}

The properties within each model stub are defined as the following:

Model Property Definition Examples
DOMAIN The type of solution the model is architected and trained for cv, nlp
SUB_DOMAIN The sub type of solution the model is architected and trained for classification, segmentation
ARCHITECTURE The name of the guiding setup for the network's graph resnet_v1, mobilenet_v1
SUB_ARCHITECTURE (optional) The scaled version of the architecture such as width or depth 50, 101, 152
FRAMEWORK The machine learning framework the model was defined and trained in pytorch, tensorflow_v1
REPO The model repository the model and baseline weights originated from sparseml, torchvision
DATASET The dataset the model was trained on imagenet, cifar10
TRAINING_SCHEME (optional) A description on how the model was trained augmented, lower_lr
SPARSE_NAME An overview of what was done to sparsify the model base, pruned, quant (quantized), pruned_quant, arch (architecture modified)
SPARSE_CATEGORY Descriptor on the degree to which the model is sparsified as compared with the baseline metric none, conservative (100% baseline), moderate (>= 99% baseline), aggressive (< 99%)
SPARSE_TARGET (optional) Descriptor for the target environment the model was sparsified for disk, edge, deepsparse, gpu

The contents of each model are made up of the following:

  • model.md: The model card containing metadata, descriptions, and information for the model.

  • model.onnx: The ONNX representation of the model’s graph.

  • model.onnx.tar.gz: A compressed format for the ONNX file. Currently ONNX does not support sparse tensors and quantized sparse tensors well for compression.

  • [FRAMEWORK]/model.[EXTENSION]: The native ML framework file(s) for the model in which it was originally trained. Such as PyTorch, Keras, TensorFlow V1

  • recipes/original.[md|yaml]: The original sparsification recipe used to create the model.

  • recipes/[NAME].[md|yaml]: Additional sparsification recipes that can be used with the model such as transfer learning.

  • sample-originals: The original sample data without any preprocessing for use with the model.

  • sample-inputs: The sample data after pre processing for use with the model.

  • sample-outputs: The outputs after running the sample inputs through the model.

  • sample-labels: The labels that classify the sample inputs.

Python APIS

The Python APIs respect this format enabling you to search and download models. Some code examples are given below.

Searching the Zoo

from sparsezoo import Zoo

models = Zoo.search_models(domain="cv", sub_domain="classification")
print(models)

Common Models

from sparsezoo.models.classification import resnet_50

model = resnet_50()
model.download()

print(model.onnx_file.downloaded_path())

Searching Optimized Versions

from sparsezoo import Zoo
from sparsezoo.models.classification import resnet_50

search_model = resnet_50()
optimized_models = Zoo.search_optimized_models(search_model)

print(optimized_models)

Console Scripts

In addition to the Python APIs, a console script entry point is installed with the package sparsezoo. This enables easy interaction straight from your console/terminal. Note, for some environments the console scripts cannot install properly. If this happens for your system and the sparsezoo command is not available, https://github.com/neuralmagic/sparsezoo/blob/main/scripts/sparsezoo.py may be used in its place.

sparsezoo -h

Searching

Search command help

sparsezoo search -h


Searching for all classification models in the computer vision domain

sparsezoo search --domain cv --sub-domain classification \
    --architecture resnet_v1 --sub-architecture 50


Searching for all ResNet-50 models

sparsezoo search --domain cv --sub-domain classification

Downloading

Download command help

sparsezoo download -h


Download ResNet-50 Model

sparsezoo download --domain cv --sub-domain classification \
    --architecture resnet_v1 --sub-architecture 50 \
    --framework pytorch --repo sparseml --dataset imagenet \
    --optim-name base --optim-category none


Download pruned and quantized ResNet-50 Model

sparsezoo download --domain cv --sub-domain classification \
    --architecture resnet_v1 --sub-architecture 50 \
    --framework pytorch --repo sparseml \
    --dataset imagenet --training-scheme augmented \
    --optim-name pruned_quant --optim-category aggressive