Achieve GPU Grade Performance on CPUs With SparseML

Deployment is the most crucial aspect of a machine learning project as most models end up being discarded either because they aren’t fast enough or are too expensive to deploy on the server. Every organization wants to build a model which is accurate and cheap for deployment. They also want to serve the user's request in a reasonable amount of time. CPU servers are cheap in comparison to their GPU counterparts but most of the time these servers are not so powerful for serving the request in a reasonable amount of time. In this blog, we would optimize a model for deploying it on a CPU server which would have performance on par with an average GPU setup.

Some Important Theories

Before we dive into the implementation part it's always good to know some theories for model optimization. We usually do the below common things for optimizing a model.

Pruning

Neural networks were created by taking inspiration from neuroscience. Their work is similar to the way our brain works. It is evident from certain research that when a child is born he has a lot more neuron connections than in his adulthood as each of these neuron connections requires nutrients and energy to be preserved hence only those connections are preserved which are required for the complete functioning of the brain.

Pruning is removing the parameters or group of parameters that have less or no contribution to prediction. Doing that saves computation and improves the performance of the model.

Quantization

In quantization, the weights of the neural network are converted from a higher precision to a lower precision like from Float32 to Float16 or Int8. This decrease the model size as well as speed up the inference.

SparseML

SparseML is a library created by the Neural Magic Team. This library provides pruning and quantization support for a wide range of models. Most of the libraries, for instance, PyTorch support unstructured pruning as a result of which inference speed up is not observed after pruning a model however that's not the case with SparseML. We can easily integrate the optimization pipeline with our PyTorch training pipeline in the upcoming sections. For inference, we use the deepspase engine which is optimized for CPU inference. SparseML has recipes for performing the pruning. Recipes are documents that have hyperparameters and a list of layers with their target pruning attributes.

Creating a Sparsified ResNet50 Model

In this section, we will use the tools from the SparseML library to create a sparsified or pruned Resnet50 model.

Installing the required libraries

          !pip install sparseml sparsezoo
!pip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html
!pip install deepsparse 

Here we start off by installing the libraries needed for optimization. SparseML for performing the pruning, sparse zoo for loading the recipe, and the pretrained model and deepsparse for inference.

Importing necessary libraries

          import sparseml
import sparsezoo
import torch
import torchvision
from sparseml.pytorch.models import ModelRegistry
from sparseml.pytorch.datasets import ImagenetteDataset, ImagenetteSize
from sparseml.pytorch.utils import ModuleExporter
from torch.utils.data import DataLoader
from torch.nn import CrossEntropyLoss
from torch.optim import SGD
from sparseml.pytorch.utils import get_prunable_layers, tensor_sparsity
from tqdm.auto import tqdm
import math
import time
import torch
from sparseml.pytorch.optim import (ScheduledModifierManager,)
from sparsezoo import Model, search_models
        

Loading the pretrained torch model and the imagenet dataset

          #Loading the torch model from Sparse zoo.
model = ModelRegistry.create(
key="resnet50",
pretrained=True,
pretrained_dataset="imagenette",
num_classes=10,)
input_shape = ModelRegistry.input_shape("resnet50")
input_size = input_shape[-1]
train_dataset = ImagenetteDataset(train=True, dataset_size=ImagenetteSize.s320, image_size=input_size)
val_dataset = ImagenetteDataset(train=False, dataset_size=ImagenetteSize.s320, image_size=input_size)

Here we are loading the pretrained SparseML Resnet50 model from the SparseModel zoo as training from scratch would take a lot of time for the imagenet dataset.

Defining the hyperparameters, loss, and the optimizer

          # setup device
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print("Using device: {}".format(device))
# setup data loaders
batch_size = 32
train_loader = DataLoader(train_dataset, batch_size, shuffle=True, pin_memory=True, num_workers=2)
val_loader = DataLoader(
val_dataset, batch_size, shuffle=False, pin_memory=True, num_workers=2)
# setup loss function and optimizer, LR will be overriden by sparseml
criterion = CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9)
        

Defining the torch training function

          def run_model_one_epoch(model, data_loader, criterion, device, train=False, optimizer=None):
 if train:
  model.train()
 else:
  model.eval()
 running_loss = 0.0
 total_correct = 0
 total_predictions = 0
 for step, (inputs, labels) in tqdm(enumerate(data_loader), total=len(data_loader)):
  inputs = inputs.to(device)
  labels = labels.to(device)
  if train:
   optimizer.zero_grad()
  outputs, _ = model(inputs)  
  loss = criterion(outputs, labels)
  if train:
   loss.backward()
   optimizer.step()
  running_loss += loss.item()
  predictions = outputs.argmax(dim=1)
  total_correct += torch.sum(predictions == labels).item()   
  total_predictions += inputs.size(0)
 loss = running_loss / (step + 1.0)
 accuracy = total_correct / total_predictions
 return loss, accuracy

Here we are defining the training function as we would do for any torch model.

          zoo_model = search_models(
domain="cv",
sub_domain="classification",
architecture="resnet_v1",
sub_architecture="50",
framework="pytorch",
repo="sparseml",
dataset="imagenette",
sparse_name="pruned",)[0]  # unwrap search result
recipe_path = zoo_model.recipes.default.path
print(f"Recipe downloaded to: {recipe_path}")
        

As mentioned in the previous section the recipes are the documents or configuration for performing the pruning. These configurations include various model training parameters for training and pruning and information like which layers to prune and how much percentage to prune so that the accuracy remains preserved.

Defining the ScheduledModifierManager and optimizer wrapper

          # create ScheduledModifierManager and Optimizer wrapper
manager = ScheduledModifierManager.from_yaml(recipe_path)
optimizer = manager.modify(model, optimizer, steps_per_epoch=len(train_loader))
# Run model pruning
epoch = manager.min_epochs
while epoch < manager.max_epochs:
 # run training loop
 epoch_name = "{}/{}".format(epoch + 1, manager.max_epochs)
 print("Running Training Epoch {}".format(epoch_name))
 train_loss, train_acc = run_model_one_epoch(model, train_loader, criterion, device, train=True, optimizer=optimizer)
 print("Training Epoch: {}\nTraining Loss: {}\nTop 1 Acc: {}\n".format(epoch_name, train_loss, train_acc))
 # run validation loop
 print("Running Validation Epoch {}".format(epoch_name))
 val_loss, val_acc = run_model_one_epoch(model, val_loader, criterion, device)
 print("Validation Epoch: {}\nVal Loss: {}\nTop 1 Acc: {}\n".format(epoch_name, val_loss, val_acc))
 epoch += 1
 manager.finalize(model)
        

Here we are defining the SparseML optimizer wrapper and wrapping our torch model in the optimizer. This wrapper would perform all the operations required for pruning or quantizing the model. Finally, we have defined the training loop where the training is performed.

Export the trained model into ONNX format

          save_dir = "pytorch_classification"
exporter = ModuleExporter(model, output_dir=save_dir)
exporter.export_pytorch(name="resnet50_imagenette_pruned.pth")exporter.export_onnx(torch.randn(1, 3, 224,224),name="resnet50_imagenette_pruned.onnx")
        

Here we are exporting the trained model to ONNX format accepted by the deepsparse engine.

Loading the trained model and doing benchmarking

          batch_size=1
#Loading pretrained torch model.
resnet_torch=resnet50(pretrained=True)
resnet_torch=resnet_torch.eval()
#Loadig the trained model.
engine = compile_model('/content/pytorch_classification/resnet50_imagenette_pruned.onnx', batch_size)
dummy_image=torch.randn((1,3,224,224))
#Generating random inputs.
inputs = generate_random_inputs('/content/content/pytorch_classification/resnet50_imagenette_pruned.onnx', batch_size)
#Warmup.
outputs = engine.run(inputs)
pred=resnet_torch(dummy_image)
#Benchmark torch model.
t1=time.time()
for _ in range(100):
 pred=resnet_torch(dummy_image)
t2=time.time()
print("Time taken torch=",(t2-t1))
#Benchmark.
t1=time.time()
for _ in range(100):
 outputs = engine.run(inputs)
t2=time.time()
print("Time taken deepsparse=",(t2-t1))

        

We can see in the above figure that the optimized model is around 2.36 times faster. This was tested on colab pro high RAM runtime which has 4-core Xeon CPUs. Furthermore, the speedup can be brought by quantizing the model to int8 as Xeon supports quantized inference.

This section was mentioned for explaining how we can use the optimization in a custom model. SparseML also has readymade pipeline integration for various models like YOLO v3, Yolo v5, BERT, etc. Let's try creating a pruned and quantized YOLO v3 model in the next section.

Creating a Sparse and Quantized YOLO v3 Model

In this section, we would train a custom YOLO v3 model using the integration provided by the SparseML team and would benchmark against the equivalent torch implementation.

Clone the SparseML integration repo from GitHub

          !git clone https://github.com/neuralmagic/sparseml.git

Here we have cloned the SparseML repo which contains the integration source code for various models.

Download and extract our stop sign dataset

          !gdown --id 1j4w1UXzCTkzkiKnwj1Luo1x0ie9Yse_b
!unzip /content/stop_sign_dataset_yolov7.zip 

Navigate to the yolov3 implementation folder and run the setup bash script

          %cd /content/sparseml/integrations/ultralytics-yolov3
!bash setup_integration.sh
        

The integration folder within the sparseML folder contains a separate folder for different models having the setup scripts and training pipelines source code for them. Hence we now navigate into it. We then run the setup_integration bash script to set up the dependencies for running the optimization code. It is a folder named yolov3 which contains the training script.

Copying the extracted data to the yolov3 folder

          !cp -r /content/train /content/sparseml/integrations/ultralytics-yolov3/yolov3/
!cp -r /content/valid /content/sparseml/integrations/ultralytics-yolov3/yolov3/
!cp  /content/data.yaml /content/sparseml/integrations/ultralytics-yolov3/yolov3/
        

Here we are moving the train test and data.yaml to the newly obtained yolov3 folder.

Rewriting the yaml file to change the path to the images.

          %%writefile data.yaml
train: ./train/images
val: ./valid/images
test: ./test/images
nc: 1
names: ['stop']
        

The train, Val, and test attributes point to the image locations of train, test, and split sets respectively. In the stop sign dataset, they were pointing to some other locations so we edit it here.

Editing the recipe file for lowering the quantization epoch

          <!--
Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->
---
# General Epoch/LR variables
num_epochs: &num_epochs 100
init_lr: &init_lr 0.0032
# Pruning Epoch/LR variables
pruning_recovery_start_lr: &pruning_recovery_start_lr 0.0032
pruning_recovery_lr_step_size: &pruning_recovery_lr_step_size 1
pruning_recovery_lr_gamma: &pruning_recovery_lr_gamma 0.985
# Quantization Epoch/LR variables
quantization_start_epoch: &quantization_start_epoch 99
quantization_init_lr: &quantization_init_lr 0.0032
# Block pruning management variables
pruning_start_epoch: &pruning_start_epoch 0
pruning_end_epoch: &pruning_end_epoch 95
pruning_update_frequency: &pruning_update_frequency 0.5
init_sparsity: &init_sparsity 0.05
mask_type: &mask_type [1, 4]
prune_none_target_sparsity: &prune_none_target_sparsity 0.4
prune_low_target_sparsity: &prune_low_target_sparsity 0.75
prune_mid_target_sparsity: &prune_mid_target_sparsity 0.8
prune_high_target_sparsity: &prune_high_target_sparsity 0.85
  
# Modifiers
training_modifiers:
  - !EpochRangeModifier
    start_epoch: 0.0
    end_epoch: *num_epochs
# pruning phase
  - !SetLearningRateModifier
    start_epoch: 0.0
    learning_rate: *init_lr
# pruning recovery phase
  - !LearningRateModifier
    constant_logging: False
    end_epoch: *quantization_start_epoch
    init_lr: *pruning_recovery_start_lr
    lr_class: StepLR
    lr_kwargs: {
        'step_size': *pruning_recovery_lr_step_size,
        'gamma': *pruning_recovery_lr_gamma
      }
    start_epoch: *pruning_end_epoch
    update_frequency: -1.0
    
  - !SetWeightDecayModifier
    start_epoch: *pruning_end_epoch
    weight_decay: 0.0
    
  # quantization aware training phase
  - !SetLearningRateModifier
    start_epoch: *quantization_start_epoch
    learning_rate: *quantization_init_lr
pruning_modifiers:
  - !GMPruningModifier
    params:
      - model.16.conv.weight
      - model.19.cv1.conv.weight
      - model.26.cv2.conv.weight
    init_sparsity: *init_sparsity
    final_sparsity: *prune_none_target_sparsity
    mask_type: *mask_type
    start_epoch: *pruning_start_epoch
    end_epoch: *pruning_end_epoch
    update_frequency: *pruning_update_frequency
- !GMPruningModifier
    params:
      - model.12.cv1.conv.weight
      - model.14.conv.weight
      - model.2.cv1.conv.weight
      - model.20.cv2.conv.weight
      - model.22.conv.weight
      - model.26.cv1.conv.weight
      - model.27.1.cv1.conv.weight
      - model.28.m.0.weight
      - model.28.m.2.weight
      - model.4.0.cv1.conv.weight
      - model.6.1.cv1.conv.weight
      - model.6.2.cv1.conv.weight
      - model.6.3.cv1.conv.weight
      - model.6.4.cv1.conv.weight
      - model.6.5.cv1.conv.weight
      - model.6.6.cv1.conv.weight
      - model.6.7.cv1.conv.weight
      - model.8.0.cv1.conv.weight
      - model.8.1.cv1.conv.weight
      - model.8.2.cv1.conv.weight
      - model.8.3.cv1.conv.weight
      - model.8.4.cv1.conv.weight
      - model.8.5.cv1.conv.weight
      - model.8.6.cv1.conv.weight
      - model.8.7.cv1.conv.weight
    init_sparsity: *init_sparsity
    final_sparsity: *prune_low_target_sparsity
    mask_type: *mask_type
    start_epoch: *pruning_start_epoch
    end_epoch: *pruning_end_epoch
    update_frequency: *pruning_update_frequency
- !GMPruningModifier
    params:
      - model.1.conv.weight
      - model.10.0.cv1.conv.weight
      - model.10.1.cv1.conv.weight
      - model.10.2.cv1.conv.weight
      - model.10.3.cv1.conv.weight
      - model.11.cv1.conv.weight
      - model.12.cv2.conv.weight
      - model.19.cv2.conv.weight
      - model.2.cv2.conv.weight
      - model.27.0.cv1.conv.weight
      - model.27.0.cv2.conv.weight
      - model.27.1.cv2.conv.weight
      - model.28.m.1.weight
      - model.3.conv.weight
      - model.4.0.cv2.conv.weight
      - model.4.1.cv1.conv.weight
      - model.4.1.cv2.conv.weight
      - model.5.conv.weight
      - model.6.0.cv1.conv.weight
      - model.6.0.cv2.conv.weight
      - model.6.1.cv2.conv.weight
      - model.6.2.cv2.conv.weight
      - model.6.3.cv2.conv.weight
    init_sparsity: *init_sparsity
    final_sparsity: *prune_mid_target_sparsity
    mask_type: *mask_type
    start_epoch: *pruning_start_epoch
    end_epoch: *pruning_end_epoch
    update_frequency: *pruning_update_frequency
- !GMPruningModifier
    params:
      - model.10.0.cv2.conv.weight
      - model.10.1.cv2.conv.weight
      - model.10.2.cv2.conv.weight
      - model.10.3.cv2.conv.weight
      - model.11.cv2.conv.weight
      - model.13.conv.weight
      - model.15.conv.weight
      - model.20.cv1.conv.weight
      - model.21.conv.weight
      - model.23.conv.weight
      - model.6.4.cv2.conv.weight
      - model.6.5.cv2.conv.weight
      - model.6.6.cv2.conv.weight
      - model.6.7.cv2.conv.weight
      - model.7.conv.weight
      - model.8.0.cv2.conv.weight
      - model.8.1.cv2.conv.weight
      - model.8.2.cv2.conv.weight
      - model.8.3.cv2.conv.weight
      - model.8.4.cv2.conv.weight
      - model.8.5.cv2.conv.weight
      - model.8.6.cv2.conv.weight
      - model.8.7.cv2.conv.weight
      - model.9.conv.weight
    init_sparsity: *init_sparsity
    final_sparsity: *prune_high_target_sparsity
    mask_type: *mask_type
    start_epoch: *pruning_start_epoch
    end_epoch: *pruning_end_epoch
    update_frequency: *pruning_update_frequency
quantization_modifiers:
  - !QuantizationModifier
    start_epoch: *quantization_start_epoch
    submodules:
      - model.0
      - model.1 
      - model.2
      - model.3 
      - model.4
      - model.5
      - model.6
      - model.7
      - model.8
      - model.9
      - model.10
      - model.11
      - model.12
      - model.13
      - model.14
      - model.15
      - model.16
      - model.17
      - model.18
      - model.19
      - model.20
      - model.21
      - model.22
      - model.23
      - model.24
      - model.25
      - model.26
      - model.27
---
# YOLOv3-SPP Pruned-Quantized
This recipe creates a sparse-quantized, [YOLOv3-SPP](https://arxiv.org/abs/1804.02767) model that achieves 94% recovery of its baseline accuracy on the COCO detection dataset.
Training was done using 4 GPUs at half precision using a total training batch size of 256 with the
[SparseML integration with ultralytics/yolov3](https://github.com/neuralmagic/sparseml/tree/main/integrations/ultralytics-yolov3).
When running, adjust hyperparameters based on training environment and dataset.
Note that half-precision, EMA, and pickling are not supported for quantization.
Therefore, once quantization is run, all three will be disabled for the training pipeline.
This additionally means that the checkpoints are saved using state_dicts rather than pickels of the model.
## Weights and Biases
- [YOLOv3-SPP LeakyReLU on VOC](https://wandb.ai/neuralmagic/yolov3-spp-lrelu-voc/runs/2dfy3rgs)
## Training
To set up the training environment, follow the instructions on the [integration README](https://github.com/neuralmagic/sparseml/blob/main/integrations/ultralytics-yolov3/README.md).
Using the given training script from the `yolov3` directory the following command can be used to launch this recipe.  
The contents of the `hyp.pruned_quantized.yaml` hyperparameters file is given below.
Adjust the script command for your GPU device setup. 
Ultralytics supports both DataParallel and DDP.
*script command:*
```
python train.py \
    --recipe ../recipes/yolov3-spp.pruned.md \
    --weights PRETRAINED_WEIGHTS \
    --data voc.yaml \
    --hyp ../data/hyp.pruned_quantized.yaml \
    --name yolov3-spp-lrelu-pruned-quantized
```
hyp.prune_quantized.yaml:
```yaml
lr0: 0.0
lrf: 0.0
momentum: 0.843
weight_decay: 0.00036
warmup_epochs: 40.0
warmup_momentum: 0.5
warmup_bias_lr: 0.05
box: 0.0296
cls: 0.243
cls_pw: 0.631
obj: 0.301
obj_pw: 0.911
iou_t: 0.2
anchor_t: 2.91
fl_gamma: 0.0
hsv_h: 0.0138
hsv_s: 0.664
hsv_v: 0.464
degrees: 0.373
translate: 0.245
scale: 0.898
shear: 0.602
perspective: 0.0
flipud: 0.00856
fliplr: 0.5
mosaic: 1.0
mixup: 0.243
```
        

Here we have edited two parameters. The &num_epochs to 100 which was 242 and &quantization_start_epoch to 99 which was 241. The reason I did this was that accuracy was not the concern of this experiment as it has been already confirmed by the SparseML team that it's nearly the same. Also, 242 was for the COCO model which has 80 classes, and hence model requires more time to recover the accuracy which is not the case with the stop sign dataset.

Navigate to the newly obtained yolov3 folder and start the training

          #Navigate to yolov3.
%cd yolov3
!python train.py \
--recipe /content/yolov3-spp.pruned_quantized.md \
--weights zoo \
--cfg /content/sparseml/integrations/ultralytics-yolov3/models/yolov3-spp.lrelu.yaml \
--data data.yaml \
--hyp data/hyp.finetune.yaml \
--epochs 100 \
--batch-size 2 \
--name yolov3-spp-leaky_relu-pruned_quant
        

Below is an explanation of the arguments above

recipe: The path to the recipe. Since we want to model to be pruned and quantized too we select yolov3-spp.pruned_quantized.md recipe

cfg: This is the model configuration. The sparseML team has replaced the default yolov3 activations with LeakyReLU for faster inference.

data: The data yaml location.

hyp: The location of the hyperparameter file.

epochs: The number of epochs to run the training.

batch_size: Batch size. The memory required for training the quantization model is usually 2-3x the original model hence be careful when choosing this.

name: Name given to the folder where the checkpoints will be stored.

Exporting the model as ONNX

          !python models/export.py --weights /content/sparseml/integrations/ultralytics-yolov3/yolov3/runs/train/yolov3-spp-leaky_relu-pruned_quant/weights/best.pt --img-size 512 512  

Finally, we export the trained model with an image size of 512x512.

Benchmarking the trained model

          onnx_filepath = "/content/content/last.onnx"
batch_size = 1
# Generate random sample inputinputs = generate_random_inputs(onnx_filepath, batch_size)
# Compile and run
engine = compile_model(onnx_filepath, batch_size)
outputs = engine.run(inputs)
t1=time.time()
for _ in range(100):
 pipeline_outputs =engine.run(inputs)
t2=time.time()
print("Time taken=",(t2-t1))
        

Benchmarking the YOLO v3-spp torch model

We need the ultralytics YOLOv3 PyTorch code to load the model hence we clone the same, navigate to the cloned directory, then load the model and finally benchmark it.

          !git clone https://github.com/ultralytics/yolov3.git
%cd yolov3
!wget https://github.com/ultralytics/yolov3/releases/download/v9.6.0/yolov3-spp.pt
model=torch.load('/content/yolov3/yolov3-spp.pt', map_location='cpu')
model=model['model'].float().eval()
t_inputs=torch.rand((1,3,512,512))
pred=model(t_inputs)
t1=time.time()
for _ in range(100):
 pred=model(t_inputs)
t2=time.time()
print("Time taken=",(t2-t1))
        

It can be observed that the pruned and quantized model runs 3.8 times faster in comparison to the torch model on a 4-core Xeon CPU. It can be further sped up to 7.6 times by selecting yolov3-spp.pruned_quantized.short.md recipe. However, it can’t be can run on the Xeon CPU as it doesn’t support the VNNI instruction set. Also, the speed difference increases with the increase in CPU cores.

Note: To reproduce the result on Google Colab Pro as of the time of publishing this blog. Select GPU runtime for training and then use high RAM runtime for inference test and on GPU the sparse model will use GPU by default

Further Improvements

* If the processor supports the VNNI instruction set we could use the VNNI-based quantization recipe which could additionally speed up the inference.

* The backbone of the network can be modified for lower-end devices.

References and Further Readings

GitHub - neuralmagic/sparseml: Libraries for applying sparsification recipes to neural networks with a few lines of code, enabling faster and smaller models

Libraries for applying sparsification recipes to neural networks with a few lines of code, enabling faster and smaller models - neuralmagic/sparseml

github.com

Link to Notebooks:

GitHub - VikasOjha666/SparseML_optimization

Contribute to VikasOjha666/SparseML_optimization development by creating an account on GitHub.

github.com

Congrats you have just learned an amazing concept in deep learning. If you liked this blog please consider sharing with your network as doing that would motivate me more to invest my time in writing blogs like this.

If you have any doubts feel free to leave them in the comment section. I will be happy to help.