Neural Magic LogoNeural Magic Logo
DeepSparse EngineSparseMLSparseZoo
Use Cases
Image Classification

Sparsifying Image Classification Models with SparseML

This page explains how to create a sparse image classification model.

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 model can be deployed with Neural Magic's DeepSparse Engine. The engine enables inference with GPU-class performance directly on your CPU.

This integration enables you to create a sparse model in two ways:

  • Sparsification of Popular Torchvision Models - easily sparsify popular torchvision image classification models.
  • Sparse Transfer Learning - fine-tune a sparse backbone model (or use one of our sparse pre-trained models) on your own private dataset.

Each option is useful in different situations:

  • Sparsification from Scratch 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 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 section requires SparseML Torchvision Install.


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 argument indicates which model to start the pruning process from. It can be a SparseZoo stub or a path to a local model.
  • --recipe_path argument instructs SparseML to run the sparsification process during the training loop. It can either be the stub of a recipe in the SparseZoo or a path to a local custom recipe. For more on creating a recipe see here.

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
>--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

The DeepSparse Engine 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 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