From 5ba3b84e9aa44498696240b7904e289b37c82d28 Mon Sep 17 00:00:00 2001 From: Mark Kurtz Date: Thu, 25 Feb 2021 11:51:19 -0500 Subject: [PATCH] Rebase release/0.1 off of main for 0.1.1 (#45) * GA code, toctree links (#32) - added tracking for docs output - added help links for docs output * Update README.md (#34) removed placeholder reference to comingsoon repo in favor of active repo * Makefile build command for nightly builds (#35) * docs updates (#36) - correcting double-slash URL issue - enhancing left nav for Help - revised mission statement provided for readme and docs * Zoo functions for loading recipes and base checkpoint weights from stubs (#37) * Update for 0.1.1 release (#40) - update python version to 0.1.1 - setup.py add in version parts and _VERSION_MAJOR_MINOR for more flexibility with dependencies between neural magic packages * base weight loading for pytorch sparse quant models (#39) * add hard check for valid recipe types in download_recipe_from_stub (#42) * Log Level Changes (#44) Logging level to warn, changed credential log to debug, increased credential duration to 30 days * Sparsification update (#43) * Update sparsification descriptions and move to preferred verbage * more sparsification updates * update from comments in deepsparse for sparsification * Update README.md Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update README.md Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update README.md Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update README.md Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update docs/source/index.rst Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update docs/source/index.rst Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update docs/source/index.rst Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update docs/source/index.rst Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update README.md Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * Update docs/source/index.rst Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> * fix links in index.rst for reviewed content * update compontent overviews and taglines * update "this repository link" for docs in overview * update doc links to not have trailing slash Co-authored-by: Jeannie Finks <74554921+jeanniefinks@users.noreply.github.com> Co-authored-by: Jeannie Finks (NM) <74554921+jeanniefinks@users.noreply.github.com> Co-authored-by: Benjamin Fineran Co-authored-by: Kevin Escobar Rodriguez --- Makefile | 3 +- README.md | 67 ++++---- docs/source/conf.py | 5 + docs/source/index.rst | 74 ++++----- docs/source/models.md | 38 ++--- docs/source/quicktour.md | 39 ++--- docs/source/recipes.md | 87 +++++----- scripts/{sparsezoo.py => sparsezoo_script.py} | 0 setup.py | 32 +++- src/sparsezoo/main.py | 9 +- src/sparsezoo/models/zoo.py | 151 +++++++++++++++++- src/sparsezoo/requests/authentication.py | 18 ++- src/sparsezoo/requests/search.py | 2 +- 13 files changed, 354 insertions(+), 171 deletions(-) rename scripts/{sparsezoo.py => sparsezoo_script.py} (100%) diff --git a/Makefile b/Makefile index 8a00c3af..3d0eeb49 100644 --- a/Makefile +++ b/Makefile @@ -7,6 +7,7 @@ DOCDIR := docs MDCHECKGLOBS := 'docs/**/*.md' 'docs/**/*.rst' 'examples/**/*.md' 'notebooks/**/*.md' 'scripts/**/*.md' MDCHECKFILES := CODE_OF_CONDUCT.md CONTRIBUTING.md DEVELOPING.md README.md +BUILD_ARGS := # set nightly to build nightly release TARGETS := "" # targets for running pytests: full,efficientnet,inception,resnet,vgg,ssd,yolo PYTEST_ARGS := "" ifneq ($(findstring full,$(TARGETS)),full) @@ -61,7 +62,7 @@ docs: # creates wheel file build: - python3 setup.py sdist bdist_wheel + python3 setup.py sdist bdist_wheel $(BUILD_ARGS) # clean package clean: diff --git a/README.md b/README.md index 391b1b61..e9ea939f 100644 --- a/README.md +++ b/README.md @@ -16,11 +16,11 @@ limitations under the License. # ![icon for SparseZoo](https://raw.githubusercontent.com/neuralmagic/sparsezoo/main/docs/source/icon-sparsezoo.png) SparseZoo -### Neural network model repository for highly sparse models and optimization recipes +### Neural network model repository for highly sparse and sparse-quantized models with matching sparsification recipes

- GitHub + GitHub Documentation @@ -44,42 +44,55 @@ limitations under the License. ## Overview -SparseZoo is a constantly-growing repository of optimized models and optimization recipes for neural networks. +SparseZoo is a constantly-growing repository of highly sparse and sparse-quantized models with matching sparsification 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](https://neuralmagic.com/blog/pruning-overview/) and [quantization](https://arxiv.org/abs/1609.07061) -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. +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. +Recipe-driven approaches built around sparsification algorithms 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. -### Related Products -- [DeepSparse](https://github.com/neuralmagic/deepsparse): CPU inference engine that delivers unprecedented performance for sparse models -- [SparseML](https://github.com/neuralmagic/sparseml): Libraries for state-of-the-art deep neural network optimization algorithms, enabling simple pipelines integration with a few lines of code -- [Sparsify](https://github.com/neuralmagic/sparsify): Easy-to-use autoML interface to optimize deep neural networks for better inference performance and a smaller footprint +## Sparsification + +Sparsification is the process of taking a trained deep learning model and removing redundant information from the overprecise and over-parameterized network resulting in a faster and smaller model. +Techniques for sparsification are all encompassing including everything from inducing sparsity using [pruning](https://neuralmagic.com/blog/pruning-overview/) and [quantization](https://arxiv.org/abs/1609.07061) to enabling naturally occurring sparsity using [activation sparsity](http://proceedings.mlr.press/v119/kurtz20a.html) or [winograd/FFT](https://arxiv.org/abs/1509.09308). +When implemented correctly, these techniques result in significantly more performant and smaller models with limited to no effect on the baseline metrics. +For example, pruning plus quantization can give over [7x improvements in performance](https://neuralmagic.com/blog/benchmark-resnet50-with-deepsparse) while recovering to nearly the same baseline accuracy. + +The Deep Sparse product suite builds on top of sparsification enabling you to easily apply the techniques to your datasets and models using recipe-driven approaches. +Recipes encode the directions for how to sparsify a model into a simple, easily editable format. +- Download a sparsification recipe and sparsified model from the [SparseZoo](https://github.com/neuralmagic/sparsezoo). +- Alternatively, create a recipe for your model using [Sparsify](https://github.com/neuralmagic/sparsify). +- Apply your recipe with only a few lines of code using [SparseML](https://github.com/neuralmagic/sparseml). +- Finally, for GPU-level performance on CPUs, deploy your sparse-quantized model with the [DeepSparse Engine](https://github.com/neuralmagic/deepsparse). + + +**Full Deep Sparse product flow:** + + ## 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}/OPTIM_NAME-OPTIM_CATEGORY-{OPTIM_TARGET}` +`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 | -| 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 | +| 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: @@ -89,8 +102,8 @@ The contents of each model are made up of the following: 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. +- 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. @@ -210,7 +223,7 @@ clone the repository and install any additional dependencies as required. ## Available Models and Recipes -A number of pre-trained baseline and optimized models across domains and sub domains are available and constantly being added. +A number of pre-trained baseline and sparsified models across domains and sub domains are available and constantly being added. For an up to date list, please consult the [available models listing](https://github.com/neuralmagic/sparsezoo/blob/main/docs/source/models.md). ## Resources and Learning More diff --git a/docs/source/conf.py b/docs/source/conf.py index 796aae08..a4ffa278 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -86,6 +86,11 @@ html_theme = "sphinx_rtd_theme" html_logo = "icon-sparsezoo.png" +html_theme_options = { + 'analytics_id': 'UA-128364174-1', # Provided by Google in your dashboard + 'analytics_anonymize_ip': False, +} + # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". diff --git a/docs/source/index.rst b/docs/source/index.rst index 2902d9a2..1918a9a6 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,7 +17,7 @@ SparseZoo |version| =================== -Neural network model repository for highly sparse models and optimization recipes +Neural network model repository for highly sparse and sparse-quantized models with matching sparsification recipes .. raw:: html @@ -48,51 +48,53 @@ Neural network model repository for highly sparse models and optimization recipe Overview ======== -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. +SparseZoo is a constantly-growing repository of highly sparse and sparse-quantized models with matching sparsification 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. +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. +Recipe-driven approaches built around sparsification algorithms 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. +`This repository `_ contains the Python API code to handle the connection and authentication to the cloud. -Related Products -================ +Sparsification +============== -- `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 +Sparsification is the process of taking a trained deep learning model and removing redundant information from the overprecise and over-parameterized network resulting in a faster and smaller model. +Techniques for sparsification are all encompassing including everything from inducing sparsity using `pruning `_ and `quantization `_ to enabling naturally occurring sparsity using `activation sparsity `_ or `winograd/FFT `_. +When implemented correctly, these techniques result in significantly more performant and smaller models with limited to no effect on the baseline metrics. +For example, pruning plus quantization can give over `7x improvements in performance `_ while recovering to nearly the same baseline accuracy. + +The Deep Sparse product suite builds on top of sparsification enabling you to easily apply the techniques to your datasets and models using recipe-driven approaches. +Recipes encode the directions for how to sparsify a model into a simple, easily editable format. +- Download a sparsification recipe and sparsified model from the `SparseZoo `_. +- Alternatively, create a recipe for your model using `Sparsify `_. +- Apply your recipe with only a few lines of code using `SparseML `_. +- Finally, for GPU-level performance on CPUs, deploy your sparse-quantized model with the `DeepSparse Engine `_. + + +**Full Deep Sparse product flow:** + + Resources and Learning More =========================== -- `SparseML Documentation `_ -- `Sparsify Documentation `_ -- `DeepSparse Documentation `_ -- `Neural Magic Blog `_, - `Resources `_, - `Website `_ +- `SparseML Documentation `_ +- `Sparsify Documentation `_ +- `DeepSparse Documentation `_ +- `Neural Magic Blog `_, + `Resources `_, + `Website `_ Release History =============== Official builds are hosted on PyPi -- stable: `sparsezoo `_ -- nightly (dev): `sparsezoo-nightly `_ +- stable: `sparsezoo `_ +- nightly (dev): `sparsezoo-nightly `_ Additionally, more information can be found via -`GitHub Releases `_. +`GitHub Releases `_. .. toctree:: :maxdepth: 3 @@ -110,9 +112,9 @@ Additionally, more information can be found via api/sparsezoo .. toctree:: - :maxdepth: 2 - :caption: Help and Support + :maxdepth: 3 + :caption: Help - `Bugs, Feature Requests `_ - `Support, General Q&A `_ - \ No newline at end of file + Bugs, Feature Requests + Support, General Q&A + Neural Magic Docs diff --git a/docs/source/models.md b/docs/source/models.md index 3ddec0ee..6de95d24 100644 --- a/docs/source/models.md +++ b/docs/source/models.md @@ -18,33 +18,34 @@ limitations under the License. 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}` +`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 | -| 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 | +| 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](https://onnx.ai/) 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. +- 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. @@ -76,14 +77,15 @@ The contents of each model are made up of the following: | cv/classification/resnet_v1-152/pytorch/sparseml/imagenet/pruned-moderate | 77.5% top1 accuracy | | cv/classification/resnet_v1-152/pytorch/torchvision/imagenet/base-none | 77.5% top1 accuracy | | cv/classification/resnet_v1-18/pytorch/sparseml/imagenet/base-none | 69.8% top1 accuracy | -| cv/classification/resnet_v1-18/pytorch/sparseml/imagenet/sparse-conservative | 69.8% top1 accuracy | +| cv/classification/resnet_v1-18/pytorch/sparseml/imagenet/pruned-conservative | 69.8% top1 accuracy | | cv/classification/resnet_v1-18/pytorch/torchvision/imagenet/base-none | 69.8% top1 accuracy | | cv/classification/resnet_v1-34/pytorch/sparseml/imagenet/base-none | 73.3% top1 accuracy | -| cv/classification/resnet_v1-34/pytorch/sparseml/imagenet/sparse-conservative | 73.3% top1 accuracy | +| cv/classification/resnet_v1-34/pytorch/sparseml/imagenet/pruned-conservative | 73.3% top1 accuracy | | cv/classification/resnet_v1-34/pytorch/torchvision/imagenet/base-none | 73.3% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/base-none | 76.1% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-conservative | 76.1% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-moderate | 75.3% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned_quant-moderate | 75.4% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenet-augmented/pruned_quant-aggressive | 76.1% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/base-none | 99.9% top1 accuracy | | cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative | 99.9% top1 accuracy | diff --git a/docs/source/quicktour.md b/docs/source/quicktour.md index 78cbc638..62025650 100644 --- a/docs/source/quicktour.md +++ b/docs/source/quicktour.md @@ -18,33 +18,34 @@ limitations under the License. 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}` +`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 | -| 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 | +| 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](https://onnx.ai/) 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. +- 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. @@ -52,8 +53,7 @@ The contents of each model are made up of the following: ### Python APIS -The Python APIs respect this format enabling you to search and download models. -Some code examples are given below. +The Python APIs respect this format enabling you to search and download models. Some code examples are given below. #### Searching the Zoo @@ -92,7 +92,8 @@ 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, `scripts/sparsezoo.py` may be used in its place. +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. ```shell script sparsezoo -h diff --git a/docs/source/recipes.md b/docs/source/recipes.md index b059773b..7f65acf3 100644 --- a/docs/source/recipes.md +++ b/docs/source/recipes.md @@ -19,55 +19,56 @@ limitations under the License. Each recipe in the SparseZoo is stored under the model created with it and 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}/RECIPE_NAME.[md|yaml]` +`DOMAIN/SUB_DOMAIN/ARCHITECTURE{-SUB_ARCHITECTURE}/FRAMEWORK/REPO/DATASET{-TRAINING_SCHEME}/SPARSE_NAME-SPARSE_CATEGORY-{SPARSE_TARGET}?recipe-type=RECIPE_TYPE` 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 | -| RECIPE_NAME | A named descriptor for the recipe signifying what the recipe is for | optimization, transfer_learning +| 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 | +| RECIPE_TYPE | A named descriptor for the recipe signifying what the recipe is for | original, transfer_learn | ### Image Classification -| Model Tag | Validation Baseline Metric | -| ------------------------------------------------------------------------------------------------------------ | -------------------------- | -| cv/classification/efficientnet-b0/pytorch/sparseml/imagenet/arch-moderate/optimization.md | 76.5% top1 accuracy | -| cv/classification/efficientnet-b4/pytorch/sparseml/imagenet/arch-moderate/optimization.md | 82.1% top1 accuracy | -| cv/classification/inception_v3/pytorch/sparseml/imagenet/pruned-conservative/optimization.md | 77.4% top1 accuracy | -| cv/classification/inception_v3/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 76.6% top1 accuracy | -| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/base-none/optimization.md | 70.9% top1 accuracy | -| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned-conservative/optimization.md | 70.9% top1 accuracy | -| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 70.1% top1 accuracy | -| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned_quant-moderate/optimization.md | 70.1% top1 accuracy | -| cv/classification/resnet_v1-101/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 76.6% top1 accuracy | -| cv/classification/resnet_v1-152/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 77.5% top1 accuracy | -| cv/classification/resnet_v1-18/pytorch/sparseml/imagenet/sparse-conservative/optimization.md | 69.8% top1 accuracy | -| cv/classification/resnet_v1-34/pytorch/sparseml/imagenet/sparse-conservative/optimization.md | 73.3% top1 accuracy | -| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-conservative/optimization.md | 76.1% top1 accuracy | -| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 75.3% top1 accuracy | -| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet-augmented/pruned_quant-aggressive/optimization.md | 76.1% top1 accuracy | -| cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative/optimization.md | 99.9% top1 accuracy | -| cv/classification/resnet_v1-50/pytorch/torchvision/imagenette/pruned-conservative/optimization.md | 99.9% top1 accuracy | -| cv/classification/vgg-11/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 68.3% top1 accuracy | -| cv/classification/vgg-16/pytorch/sparseml/imagenet/pruned-conservative/optimization.md | 71.6% top1 accuracy | -| cv/classification/vgg-16/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 70.8% top1 accuracy | -| cv/classification/vgg-19/pytorch/sparseml/imagenet/pruned-moderate/optimization.md | 71.7% top1 accuracy | +| Model Tag | Validation Baseline Metric | +| ----------------------------------------------------------------------------------------------------------------- | -------------------------- | +| cv/classification/efficientnet-b0/pytorch/sparseml/imagenet/arch-moderate?recipe_type=original | 76.5% top1 accuracy | +| cv/classification/efficientnet-b4/pytorch/sparseml/imagenet/arch-moderate?recipe_type=original | 82.1% top1 accuracy | +| cv/classification/inception_v3/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 77.4% top1 accuracy | +| cv/classification/inception_v3/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 76.6% top1 accuracy | +| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/base-none?recipe_type=original | 70.9% top1 accuracy | +| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 70.9% top1 accuracy | +| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 70.1% top1 accuracy | +| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned_quant-moderate?recipe_type=original | 70.1% top1 accuracy | +| cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned_quant-moderate?recipe_type=original | 70.1% top1 accuracy | +| cv/classification/resnet_v1-101/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 76.6% top1 accuracy | +| cv/classification/resnet_v1-152/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 77.5% top1 accuracy | +| cv/classification/resnet_v1-18/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 69.8% top1 accuracy | +| cv/classification/resnet_v1-34/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 73.3% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 76.1% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 75.3% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/sparseml/imagenet-augmented/pruned_quant-aggressive?recipe_type=original | 76.1% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative?recipe_type=original | 99.9% top1 accuracy | +| cv/classification/resnet_v1-50/pytorch/torchvision/imagenette/pruned-conservative?recipe_type=original | 99.9% top1 accuracy | +| cv/classification/vgg-11/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 68.3% top1 accuracy | +| cv/classification/vgg-16/pytorch/sparseml/imagenet/pruned-conservative?recipe_type=original | 71.6% top1 accuracy | +| cv/classification/vgg-16/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 70.8% top1 accuracy | +| cv/classification/vgg-19/pytorch/sparseml/imagenet/pruned-moderate?recipe_type=original | 71.7% top1 accuracy | ### Object Detection -| Model Tag | Validation Baseline Metric | -| ------------------------------------------------------------------------------------------------------------ | -------------------------- | -| cv/detection/ssd-resnet50_300/pytorch/sparseml/coco/pruned-moderate/optimization.md | 41.8 mAP@0.5 | -| cv/detection/ssd-resnet50_300/pytorch/sparseml/voc/pruned-moderate/optimization.md | 51.5 mAP@0.5 | -| cv/detection/yolo_v3-spp/pytorch/ultralytics/coco/pruned-aggressive/optimization.md | 62.1 mAP@0.5 | +| Model Tag | Validation Baseline Metric | +| ----------------------------------------------------------------------------------------------------------------- | -------------------------- | +| cv/detection/ssd-resnet50_300/pytorch/sparseml/coco/pruned-moderate?recipe_type=original | 41.8 mAP@0.5 | +| cv/detection/ssd-resnet50_300/pytorch/sparseml/voc/pruned-moderate?recipe_type=original | 51.5 mAP@0.5 | +| cv/detection/yolo_v3-spp/pytorch/ultralytics/coco/pruned-aggressive?recipe_type=original | 62.1 mAP@0.5 | diff --git a/scripts/sparsezoo.py b/scripts/sparsezoo_script.py similarity index 100% rename from scripts/sparsezoo.py rename to scripts/sparsezoo_script.py diff --git a/setup.py b/setup.py index c464a55d..4c108957 100644 --- a/setup.py +++ b/setup.py @@ -12,9 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import sys +from datetime import date from typing import Tuple, List, Dict from setuptools import find_packages, setup +_PACKAGE_NAME = "sparsezoo" +_VERSION = "0.1.1" +_VERSION_MAJOR, _VERSION_MINOR, _VERSION_BUG = _VERSION.split(".") +_VERSION_MAJOR_MINOR = f"{_VERSION_MAJOR}.{_VERSION_MINOR}" +_NIGHTLY = "nightly" in sys.argv + +if _NIGHTLY: + _PACKAGE_NAME += "-nightly" + _VERSION += "." + date.today().strftime("%Y%m%d") + # remove nightly param so it does not break bdist_wheel + sys.argv.remove("nightly") + _deps = [ "numpy>=1.0.0", "onnx>=1.0.0,<1.8.0", @@ -65,17 +79,21 @@ def _setup_long_description() -> Tuple[str, str]: setup( - name="sparsezoo", - version="0.1.0", + name=_PACKAGE_NAME, + version=_VERSION, author="Neuralmagic, Inc.", author_email="support@neuralmagic.com", - description="Neural network model repository for highly sparse models " - "and optimization recipes.", + description=( + "Neural network model repository for highly sparse and sparse-quantized " + "models with matching sparsification recipes" + ), long_description=_setup_long_description()[0], long_description_content_type=_setup_long_description()[1], - keywords="inference, machine learning, neural network, deep learning model, " - "models, computer vision, nlp, pretrained transfer learning, sparsity, pruning, " - "quantization, sparse models, resnet, mobilenet, yolov3", + keywords=( + "inference, machine learning, neural network, deep learning model, models, " + "computer vision, nlp, pretrained transfer learning, sparsity, pruning, " + "quantization, sparse models, resnet, mobilenet, yolov3" + ), license="Apache", url="https://github.com/neuralmagic/sparsezoo", package_dir={"": "src"}, diff --git a/src/sparsezoo/main.py b/src/sparsezoo/main.py index a4f9a814..8ac41992 100644 --- a/src/sparsezoo/main.py +++ b/src/sparsezoo/main.py @@ -179,7 +179,6 @@ DOWNLOAD_COMMAND = "download" SEARCH_COMMAND = "search" -logging.basicConfig(level=logging.INFO) LOGGER = logging.getLogger() @@ -330,13 +329,7 @@ def _get_command_from_model(model: Model): if hasattr(model, field) and getattr(model, field) is not None ] - command_string = download_command + " ".join( - # [ - # " \\ \n ".join(command_strings[i : i + 2]) - # for i in range(0, len(command_strings), 2) - # ] - command_strings - ) + command_string = download_command + " ".join(command_strings) return command_string diff --git a/src/sparsezoo/models/zoo.py b/src/sparsezoo/models/zoo.py index e6211220..6835d74c 100644 --- a/src/sparsezoo/models/zoo.py +++ b/src/sparsezoo/models/zoo.py @@ -17,14 +17,62 @@ """ -from typing import List, Union +import warnings +from typing import Dict, List, Tuple, Union from sparsezoo.objects.model import Model -from sparsezoo.objects.optimization_recipe import OptimizationRecipe +from sparsezoo.objects.optimization_recipe import ( + OptimizationRecipe, + OptimizationRecipeTypes, +) from sparsezoo.requests import ModelArgs, download_get_request, search_get_request -__all__ = ["Zoo"] +__all__ = [ + "ZOO_STUB_PREFIX", + "parse_zoo_stub", + "Zoo", +] + + +# optional prefix for stubs +ZOO_STUB_PREFIX = "zoo:" + + +def parse_zoo_stub( + stub: str, valid_params: Union[List[str], None] = None +) -> Tuple[str, Dict[str, str]]: + """ + :param stub: A SparseZoo model stub. i.e. 'model/stub/path', + 'zoo:model/stub/path', 'zoo:model/stub/path?param1=value1¶m2=value2' + :param valid_params: list of expected parameter names to be encoded in the + stub. Will raise a warning if any unexpected param names are given. Leave + as None to not raise any warnings. Default is None + :return: the parsed base stub and a dictionary of parameter names and their values + """ + # strip optional zoo stub prefix + if stub.startswith(ZOO_STUB_PREFIX): + stub = stub[len(ZOO_STUB_PREFIX) :] + + if "?" not in stub: + return stub, {} + + stub_parts = stub.split("?") + if len(stub_parts) > 2: + raise ValueError( + "Invalid SparseZoo stub, query string must be preceded by only one '?'" + f"given {stub}" + ) + stub, params = stub_parts + params = dict(param.split("=") for param in params.split("&")) + + if valid_params is not None and any(param not in valid_params for param in params): + warnings.warn( + f"Invalid query string for stub {stub} valid params include {valid_params}," + f" given {list(params.keys())}" + ) + + return stub, params class Zoo: @@ -119,6 +167,9 @@ def load_model_from_stub( :param force_token_refresh: True to refresh the auth token, False otherwise :return: The requested Model instance """ + if isinstance(stub, str): + stub, _ = parse_zoo_stub(stub, valid_params=[]) + response_json = download_get_request( args=stub, file_name=None, @@ -443,3 +494,97 @@ def search_optimized_recipes( match_training_scheme=match_training_scheme, ) return [recipe for model in optimized_models for recipe in model.recipes] + + @staticmethod + def download_recipe_from_stub( + stub: str, + ) -> str: + """ + :param stub: a string model stub that points to a SparseZoo model. + recipe_type may be added as a stub parameter. i.e. + "model/stub/path", "zoo:model/stub/path", + "zoo:model/stub/path?recipe_type=original" + :return: file path of the downloaded recipe for that model + """ + stub, args = parse_zoo_stub(stub, valid_params=["recipe_type"]) + recipe_type = _get_stub_args_recipe_type(args) + model = Zoo.load_model_from_stub(stub) + + for recipe in model.recipes: + if recipe.recipe_type == recipe_type: + return recipe.downloaded_path() + + found_recipe_types = [recipe.recipe_type for recipe in model.recipes] + raise RuntimeError( + f"No recipe with recipe_type {recipe_type} found for model {model}. " + f"Found {len(model.recipes)} recipes with recipe types {found_recipe_types}" + ) + + @staticmethod + def download_recipe_base_framework_files( + stub: str, + extensions: Union[List[str], None] = None, + ) -> List[str]: + """ + :param stub: a string model stub that points to a SparseZoo model. + recipe_type may be added as a stub parameter. i.e. + "model/stub/path", "zoo:model/stub/path", + "zoo:model/stub/path?recipe_type=transfer" + :param extensions: List of file extensions to filter for. ex ['.pth', '.ptc']. + If None or empty list, all framework files are downloaded. Default is None + :return: file path to the downloaded framework checkpoint files for the + base weights of this recipe + """ + stub, args = parse_zoo_stub(stub, valid_params=["recipe_type"]) + recipe_type = _get_stub_args_recipe_type(args) + model = Zoo.load_model_from_stub(stub) + + if recipe_type == OptimizationRecipeTypes.TRANSFER_LEARN.value: + # return final model's optimized weights for sparse transfer learning + framework_files = model.download_framework_files(extensions=extensions) + + # download only pre-quantized weights if available + checkpoint_framework_files = [ + framework_file + for framework_file in framework_files + if ".ckpt" in framework_file + ] + + # return non-empty list, preferring filtered list + return checkpoint_framework_files or framework_files + else: + # search for base model, and return those weights as a starting checkpoint + base_model = [ + result + for result in Zoo.search_optimized_models(model) + if result.optim_name == "base" + ] + if not base_model: + raise ValueError(f"Could not find base model for model {model}") + framework_files = base_model[0].download_framework_files( + extensions=extensions + ) + + # filter out checkpoint weights if any exist + base_framework_files = [ + framework_file + for framework_file in framework_files + if ".ckpt" not in framework_file + ] + + # return non-empty list, preferring filtered list + return base_framework_files or framework_files + + +def _get_stub_args_recipe_type(stub_args: Dict[str, str]) -> str: + # check recipe type, default to original, and validate + recipe_type = stub_args.get("recipe_type", OptimizationRecipeTypes.ORIGINAL.value) + + # validate + valid_recipe_types = list(map(lambda typ: typ.value, OptimizationRecipeTypes)) + if recipe_type not in valid_recipe_types: + raise ValueError( + f"Invalid recipe_type: '{recipe_type}'. " + f"Valid recipe types: {valid_recipe_types}" + ) + return recipe_type diff --git a/src/sparsezoo/requests/authentication.py b/src/sparsezoo/requests/authentication.py index 9b56a8c8..60d83db1 100644 --- a/src/sparsezoo/requests/authentication.py +++ b/src/sparsezoo/requests/authentication.py @@ -54,7 +54,7 @@ class SparseZooCredentials: def __init__(self): if os.path.exists(CREDENTIALS_YAML): - _LOGGER.info(f"Loading sparse zoo credentials from {CREDENTIALS_YAML}") + _LOGGER.debug(f"Loading sparse zoo credentials from {CREDENTIALS_YAML}") with open(CREDENTIALS_YAML) as credentials_file: credentials_yaml = yaml.safe_load(credentials_file) if credentials_yaml and CREDENTIALS_YAML_TOKEN_KEY in credentials_yaml: @@ -66,7 +66,9 @@ def __init__(self): self._token = None self._created = None else: - _LOGGER.info(f"No sparse zoo credentials files found at {CREDENTIALS_YAML}") + _LOGGER.debug( + f"No sparse zoo credentials files found at {CREDENTIALS_YAML}" + ) self._token = None self._created = None @@ -78,7 +80,7 @@ def save_token(self, token: str, created: float): :param token: the jwt for accessing sparse zoo APIs :param created: the approximate time the token was created """ - _LOGGER.info(f"Saving sparse zoo credentials at {CREDENTIALS_YAML}") + _LOGGER.debug(f"Saving sparse zoo credentials at {CREDENTIALS_YAML}") if not os.path.exists(CREDENTIALS_YAML): create_parent_dirs(CREDENTIALS_YAML) with open(CREDENTIALS_YAML, "w+") as credentials_file: @@ -99,17 +101,17 @@ def token(self): """ :return: obtain the token if under 1 day old, else return None """ - _LOGGER.info(f"Obtaining sparse zoo credentials from {CREDENTIALS_YAML}") + _LOGGER.debug(f"Obtaining sparse zoo credentials from {CREDENTIALS_YAML}") if self._token and self._created is not None: creation_date = datetime.fromtimestamp(self._created, tz=timezone.utc) creation_difference = datetime.now(tz=timezone.utc) - creation_date - if creation_difference.days == 0: + if creation_difference.days < 30: return self._token else: - _LOGGER.warning(f"Expired sparse zoo credentials at {CREDENTIALS_YAML}") + _LOGGER.debug(f"Expired sparse zoo credentials at {CREDENTIALS_YAML}") return None else: - _LOGGER.warning(f"No sparse zoo credentials found at {CREDENTIALS_YAML}") + _LOGGER.debug(f"No sparse zoo credentials found at {CREDENTIALS_YAML}") return None @@ -138,7 +140,7 @@ def get_auth_header( if token and not force_token_refresh: return {NM_TOKEN_HEADER: token} elif authentication_type.lower() == PUBLIC_AUTH_TYPE: - _LOGGER.warning("Obtaining new sparse zoo credentials token") + _LOGGER.info("Obtaining new sparse zoo credentials token") created = time.time() response = requests.post( url=AUTH_API, data=json.dumps({"authentication_type": PUBLIC_AUTH_TYPE}) diff --git a/src/sparsezoo/requests/search.py b/src/sparsezoo/requests/search.py index 084f5444..83c22f46 100644 --- a/src/sparsezoo/requests/search.py +++ b/src/sparsezoo/requests/search.py @@ -64,7 +64,7 @@ def search_get_request( search_args = "&".join(search_args) url = f"{BASE_API_URL}/{SEARCH_PATH}/{args.model_url_root}?{search_args}" - _LOGGER.debug(f"Searching objects from {url}") + _LOGGER.info(f"Searching objects from {url}") response_json = requests.get(url=url, headers=header).json() return response_json