SparseML Image Classification pipeline integrates with torch and torchvision libraries to enable the sparsification of popular image classification model. Sparsification is a powerful technique that results in faster, smaller, and cheaper deployable models. After training, the sparse model can be deployed with DeepSparse for GPU-class performance directly on your CPU.

This integration enables you to create a sparse model in two ways. Each option is useful in different situations:

Sparsification of Popular Torchvision Models—Easily sparsify popular torchvision image classification models. This enables you to create a sparse version of any model (even those not in the SparseZoo), but requires hand-tuning the hyperparameters of the Sparsification algorithm.
Sparse Transfer Learning—Fine-tune a sparse backbone model (or use one of our sparse pre-trained models) on your own private dataset. This is the easiest path to creating a sparse model trained on your data. Simply pull a pre-sparsified model and transfer learning recipe from the SparseZoo and fine-tune on your data with a single command.

Installation Requirements

This use case requires installation of SparseML Torchvision.

Tutorials

Here are additional tutorials for this functionality:

Getting Started

Sparsifying Image Classification Models

In the example below, a dense ResNet model is sparsified and fine-tuned on the Imagenette dataset.

1sparseml.image_classification.train \
2    --recipe-path "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative?recipe_type=original" \
3    --dataset-path ./data \
4    --pretrained True \
5    --arch-key resnet50 \
6    --dataset imagenette \
7    --train-batch-size 128 \
8    --test-batch-size 256 \
9    --loader-num-workers 8 \
10    --save-dir sparsification_example \
11    --logs-dir sparsification_example \
12    --model-tag resnet50-imagenette-pruned \
13    --save-best-after 8

The most important arguments are --dataset_path and --recipe_path:

--dataset_path indicates the model with which to start the pruning process. It can be a SparseZoo stub or a path to a local model.
--recipe_path instructs SparseML to run the sparsification process during the training loop. It can be either the stub of a recipe in the SparseZoo or a path to a local custom recipe. See Creating Sparsification Recipes for more information.

Sparse Transfer Learning

SparseML also enables you to fine-tune a pre-sparsified model onto your own dataset. While you are free to use your backbone, we encourage you to leverage one of our sparse pre-trained models to boost your productivity!

The command below fetches a pruned ResNet model, pre-trained on ImageNet dataset from the SparseZoo and then fine-tunes the model on the Imagenette dataset while preserving sparsity.

1sparseml.image_classification.train \
2    --recipe-path zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned95-none?recipe_type=transfer-classification \
3    --checkpoint-path zoo \
4    --arch-key resnet50 \
5    --model-kwargs '{"ignore_error_tensors": ["classifier.fc.weight", "classifier.fc.bias"]}' \
6    --dataset imagenette \
7    --dataset-path /PATH/TO/IMAGENETTE  \
8    --train-batch-size 32 \
9    --test-batch-size 64 \
10    --loader-num-workers 0 \
11    --optim Adam \
12    --optim-args '{}' \
13    --model-tag resnet50-imagenette-transfer-learned

SparseML CLI

SparseML installation provides a CLI for sparsifying your models for a specific task. Appending the --help argument displays a full list of options for training in SparseML:

1sparseml.image_classification.train --help
>Usage: sparseml.image_classification.train [OPTIONS]
>
>PyTorch training integration with SparseML for image classification models
>
>Options:
>--train-batch-size, --train_batch_size INTEGER
>Train batch size  [required]
>--test-batch-size, --test_batch_size INTEGER
>Test/Validation batch size  [required]
>--dataset TEXT                  The dataset to use for training, ex:
>`imagenet`, `imagenette`, `cifar10`, etc.
>Set to `imagefolder` for a generic dataset
>setup with imagefolder type structure like
>imagenet or loadable by a dataset in
>`sparseml.pytorch.datasets`  [required]
>--dataset-path, --dataset_path DIRECTORY
>The root dir path where the dataset is
>stored or should be downloaded to if
>available  [required]
>--arch_key, --arch-key TEXT     The architecture key for image
>classification model; example: `resnet50`,
>`mobilenet`. Note: Will be read from the
>checkpoint if not specified
>--checkpoint-path, --checkpoint_path TEXT
>A path to a previous checkpoint to load the
>state from and resume the state for. If
>provided, pretrained will be ignored . If
>using a SparseZoo recipe, can also provide
>'zoo' to load the base weights associated
>with that recipe. Additionally, can also
>provide a SparseZoo model stub to load model
>weights from SparseZoo
>...

Exporting to ONNX

The artifacts of the training process are saved to --save-dir under --model-tag. Once the script terminates, you should find everything required to deploy or further modify the model, including the recipe (with the full description of the sparsification attributes), checkpoint files (saved in the appropriate framework format), etc.

Exporting the Sparse Model to ONNX

DeepSparse uses the ONNX format to load neural networks and then deliver breakthrough performance for CPUs by leveraging the sparsity and quantization within a network.

The SparseML installation provides a sparseml.image_classification.export_onnx command that you can use to load the checkpoint and create a new model.onnx file in the same directory where the framework directory is stored. Be sure the --model_path argument points to your trained model.pth or checkpoint-best.pth file. Both are included in <save-dir>/<model-tag>/framework/ from the sparsification run.

1sparseml.image_classification.export_onnx \
2    --arch-key resnet50 \
3    --dataset imagenet \
4    --dataset-path ./data/imagenette-160 \
5    --checkpoint-path sparsification_example/resnet50-imagenette-pruned/training/model.pth

NLP Deployments with DeepSparse

Image Classification Deployments With DeepSparse