SparseZoo is a constantly-growing repository of optimized models and optimization recipes for neural networks. It simplifies and accelerates your time-to-value in building performant deep learning models with a collection of inference-optimized models and recipes to prototype from.
Available via API and hosted in the cloud, the SparseZoo contains both baseline and models optimized to different degrees of inference performance vs baseline loss recovery. Optimizations on neural networks include approaches such as pruning and quantization allowing for significantly faster models with limited to no effect on their baseline metrics such as accuracy. Recipe-driven approaches built around these optimizations allow you to take the models as given, transfer learn from the models onto private datasets, or transfer the recipes to your architectures.
This repository contains the Python API code to handle the connection and authentication to the cloud.
- DeepSparse: CPU inference engine that delivers unprecedented performance for sparse models
- SparseML: Libraries for state-of-the-art deep neural network optimization algorithms, enabling simple pipelines integration with a few lines of code
- Sparsify: Easy-to-use autoML interface to optimize deep neural networks for better inference performance and a smaller footprint
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}/OPTIM_NAME-OPTIM_CATEGORY-{OPTIM_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 |
| OPTIM_NAME | An overview of what was done to optimize the model | base, pruned, quant (quantized), pruned_quant, arch (architecture modified) |
| OPTIM_CATEGORY | Descriptor on the degree to which the model is optimized as compared with the baseline metric | none, conservative (100% baseline), moderate (>= 99% baseline), aggressive (< 99%) |
| OPTIM_TARGET | (optional) Descriptor for the target environment the model was optimized 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/optimization.[md|yaml]: The original optimization recipe used to create the model.
- recipes/[NAME].[md|yaml]: Additional optimization 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.
The Python APIs respect this format enabling you to search and download models. Some code examples are given below.
from sparsezoo import Zoo
models = Zoo.search_models(domain="cv", sub_domain="classification")
print(models)from sparsezoo.models.classification import resnet_50
model = resnet_50()
model.download()
print(model.onnx_file.downloaded_path())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)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 -hSearch 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 classificationDownload 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 aggressiveFor a more in-depth read, check out SparseZoo documentation.
This repository is tested on Python 3.6+, and Linux/Debian systems. It is recommended to install in a virtual environment to keep your system in order.
Install with pip using:
pip install sparsezooThen if you would like to explore any of the scripts or notebooks clone the repository and install any additional dependencies as required.
A number of pre-trained baseline and optimized models across domains and sub domains are available and constantly being added. For an up to date list, please consult the available models listing.
- SparseZoo Documentation
- SparseML Documentation
- Sparsify Documentation
- DeepSparse Documentation
- Neural Magic Blog, Resources, Website
We appreciate contributions to the code, examples, and documentation as well as bug reports and feature requests! Learn how here.
For user help or questions about SparseZoo, use our GitHub Discussions. Everyone is welcome!
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 email us at [email protected] or fill out this form.
The project is licensed under the Apache License Version 2.0.
Official builds are hosted on PyPi
- stable: sparsezoo
- nightly (dev): sparsezoo-nightly
Additionally, more information can be found via GitHub Releases.