SparseZoo

Neural network model repository for highly sparse and sparse-quantized models with matching sparsification recipes

SparseZoo is a constantly-growing repository of sparsified (pruned and pruned-quantized) models with matching sparsification recipes for neural networks. SparseZoo simplifies and accelerates your time-to-value in building performant deep learning models with a collection of inference-optimized models and recipes from which to prototype. Read more about sparsification.

Available via API and hosted in the cloud, the SparseZoo contains both baseline and models sparsified to different degrees of inference performance versus baseline loss recovery. Recipe-driven approaches built around sparsification algorithms allow you to use the models as given, transfer-learn from the models onto private datasets, or transfer the recipes to your architectures.

The GitHub repository contains the Python API code to handle the connection and authentication to the cloud.

Highlights

• Model Stub Architecture Overview
• Available Model Recipes
• sparsezoo.neuralmagic.com

Installation

See the SparseZoo Installation page for installation instructions.

Quick Tour

The SparseZoo Python API enables you to search and download sparsified models. Code examples are given below. We encourage users to load SparseZoo models by copying a stub directly from a model page.

Introduction to Model Class Object

The Model is a fundamental object that serves as a main interface with the SparseZoo library. It represents a SparseZoo model, together with all its directories and files.

Creating a Model Class Object From SparseZoo Stub

1from sparsezoo import Model
2
3stub = "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95_quant-none"
4
5model = Model(stub)
6print(str(model))
>> Model(stub=zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95_quant-none)

Creating a Model Class Object From Local Model Directory

1from sparsezoo import Model
2
3directory = ".../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0"
4
5model = Model(directory)
6print(str(model))
>> Model(directory=.../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0)

Manually Specifying the Model Download Path

Unless specified otherwise, the model created from the SparseZoo stub is saved to the local SparseZoo cache directory. This can be overridden by passing the optional download_path argument to the constructor:

1from sparsezoo import Model
2
3stub = "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95_quant-none"
4download_directory = "./model_download_directory"
5
6model = Model(stub, download_path = download_directory)

Downloading the Model Files

Once the model is initialized from a stub, it may be downloaded either by calling the download() method or by invoking a path property. Both pathways are universal for all the files in SparseZoo. Invoking the path property will always trigger file download unless the file has already been downloaded.

1# method 1
2model.download()
3
4# method 2
5model_path = model.path

Inspecting the Contents of the SparseZoo Model

We call the available_files method to inspect which files are present in the SparseZoo model. Then, we select a file by calling the appropriate attribute:

1model.available_files
>> {'training': Directory(name=training),
>> 'deployment': Directory(name=deployment),
>> 'sample_inputs': Directory(name=sample_inputs.tar.gz),
>> 'sample_outputs': {'framework': Directory(name=sample_outputs.tar.gz)},
>> 'sample_labels': Directory(name=sample_labels.tar.gz),
>> 'model_card': File(name=model.md),
>> 'recipes': Directory(name=recipe),
>> 'onnx_model': File(name=model.onnx)}

We might take a closer look at the contents of the SparseZoo model:

1model_card = model.model_card
2print(model_card)
>> File(name=model.md)

1model_card_path = model.model_card.path
2print(model_card_path)
>> .../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0/model.md

Model, Directory, and File

In general, every file in the SparseZoo model shares a set of attributes: name, path, URL, and parent:

• name serves as an identifier of the file/directory.
• path points to the location of the file/directory.
• URL specifies the server address of the file/directory in question.
• parent points to the location of the parent directory of the file/directory in question.

A directory is a unique type of file that contains other files. For that reason, it has an additional files attribute.

1print(model.onnx_model)
>> File(name=model.onnx)
1print(f"File name: {model.onnx_model.name}\n"
2      f"File path: {model.onnx_model.path}\n"
3      f"File URL: {model.onnx_model.url}\n"
4      f"Parent directory: {model.onnx_model.parent_directory}")
>> File name: model.onnx
>> File path: .../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0/model.onnx
>> File URL: https://models.neuralmagic.com/cv-classification/...
>> Parent directory: .../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0

1print(model.recipes)
>> Directory(name=recipe)
1print(f"File name: {model.recipes.name}\n"
2      f"Contains: {[file.name for file in model.recipes.files]}\n"
3      f"File path: {model.recipes.path}\n"
4      f"File URL: {model.recipes.url}\n"
5      f"Parent directory: {model.recipes.parent_directory}")
>> File name: recipe
>> Contains: ['recipe_original.md', 'recipe_transfer-classification.md']
>> File path: /home/user/.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0/recipe
>> File URL: None
>> Parent directory: /home/user/.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0

Selecting Checkpoint-Specific Data

A SparseZoo model may contain several checkpoints. The model may contain a checkpoint that had been saved before the model was quantized. That checkpoint would be used for transfer learning. Another checkpoint might have been saved after the quantization step. That one usually is used directly for inference.

The recipes may also vary depending on the use case. We may want to access a recipe that was used to sparsify the dense model (recipe_original) or the one that enables us to sparse transfer learn from the already sparsified model (recipe_transfer).

There are three ways to access those specific files:

• Accessing recipes (through Python API)
• Accessing checkpoints (through Python API)
• Accessing recipies (through stub string arguments)

Accessing Recipes (Through Python API)

1available_recipes = model.recipes.available
2print(available_recipes)
>> ['original', 'transfer-classification']
1transfer_recipe = model.recipes["transfer-classification"]
2print(transfer_recipe)
>> File(name=recipe_transfer-classification.md)
1original_recipe = model.recipes.default # recipe defaults to `original`
2original_recipe_path = original_recipe.path # downloads the recipe and returns its path
3print(original_recipe_path)
>> .../.cache/sparsezoo/eb977dae-2454-471b-9870-4cf38074acf0/recipe/recipe_original.md

Accessing Checkpoints (Through Python API)

In general, we are expecting the following checkpoints to be included in the model:

• checkpoint_prepruning
• checkpoint_postpruning
• checkpoint_preqat
• checkpoint_postqat

The checkpoint that the model defaults to is the preqat state (just before the quantization step).

1from sparsezoo import Model
2
3stub = "zoo:nlp/question_answering/bert-base/pytorch/huggingface/squad/pruned_quant_3layers-aggressive_84"
4
5model = Model(stub)
6available_checkpoints = model.training.available
7print(available_checkpoints)
>> ['preqat']
1preqat_checkpoint = model.training.default # recipe defaults to `preqat`
2preqat_checkpoint_path = preqat_checkpoint.path # downloads the checkpoint and returns its path
3print(preqat_checkpoint_path)
>> .../.cache/sparsezoo/0857c6f2-13c1-43c9-8db8-8f89a548dccd/training
1[print(file.name) for file in preqat_checkpoint.files]
>> vocab.txt
>> special_tokens_map.json
>> pytorch_model.bin
>> config.json
>> training_args.bin
>> tokenizer_config.json
>> trainer_state.json
>> tokenizer.json

Accessing Recipes (Through Stub String Arguments)

You can also directly request a specific recipe/checkpoint type by appending the appropriate URL query arguments to the stub:

1from sparsezoo import Model
2
3stub = "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95_quant-none?recipe=transfer"
4
5model = Model(stub)
6
7# Inspect which files are present.
8# Note that the available recipes are restricted
9# according to the specified URL query arguments
10print(model.recipes.available)
>> ['transfer-classification']
1transfer_recipe = model.recipes.default # Now the recipes default to the one selected by the stub string arguments
2print(transfer_recipe)
>> File(name=recipe_transfer-classification.md)

Accessing Sample Data

You may easily request a sample batch of data that represents the inputs and outputs of the model.

1sample_data = model.sample_batch(batch_size = 10)
2
3print(sample_data['sample_inputs'][0].shape)
>> (10, 3, 224, 224) # (batch_size, num_channels, image_dim, image_dim)
1print(sample_data['sample_outputs'][0].shape)
>> (10, 1000) # (batch_size, num_classes)

Model Search

The function search_models enables you to quickly filter the contents of the SparseZoo repository to find the stubs of interest:

1from sparsezoo import search_models
2
3args = {
4    "domain": "cv",
5    "sub_domain": "segmentation",
6    "architecture": "yolact",
7}
8
9models = search_models(**args)
10[print(model) for model in models]
>> Model(stub=zoo:cv/segmentation/yolact-darknet53/pytorch/dbolya/coco/pruned82_quant-none)
>> Model(stub=zoo:cv/segmentation/yolact-darknet53/pytorch/dbolya/coco/pruned90-none)
>> Model(stub=zoo:cv/segmentation/yolact-darknet53/pytorch/dbolya/coco/base-none)

Environmental Variables

You can specify the directory where models and required credentials will be saved (temporarily during download) in your working machine. SPARSEZOO_MODELS_PATH is the path where the downloaded models will be saved temporarily. Default ~/.cache/sparsezoo/ SPARSEZOO_CREDENTIALS_PATH is the path where credentials.yaml will be saved. The default is ~/.cache/sparsezoo/.

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.

Downloading

Download command help:

sparsezoo.download -h

Download ResNet-50 model:

sparsezoo.download zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/base-none

Download pruned and quantized ResNet-50 model:

sparsezoo.download zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate

Searching

Search command help:

sparsezoo search -h

Search for all classification MobileNetV1 models in the computer vision domain:

sparsezoo search --domain cv --sub-domain classification --architecture mobilenet_v1

Search for all ResNet-50 models:

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

Resources

Learning More

• Documentation: SparseML, SparseZoo, Sparsify, DeepSparse
• Neural Magic: Blog, Resources

Release History

Official builds are hosted on PyPI

• Stable: sparsezoo
• Nightly (dev): sparsezoo-nightly

Additionally, more information can be found via GitHub Releases.

License

The project is licensed under the Apache License Version 2.0.

Community

Contribute

We appreciate contributions to the code, examples, integrations, and documentation as well as bug reports and feature requests! Learn how here.

Join

For user help or questions about SparseZoo, sign up or log into our Neural Magic Community Slack. We are growing the community member by member and happy to see you there. Bugs, feature requests, or additional questions can also be posted to our GitHub Issue Queue.

You can get the latest news, webinar and event invites, research papers, and other ML Performance tidbits by subscribing to the Neural Magic community.

For more general questions about Neural Magic, please fill out this form.