|
--- |
|
license: mit |
|
datasets: |
|
- imagenet-1k |
|
pipeline_tag: image-classification |
|
tags: |
|
- sparsity |
|
- vision-transformer |
|
- pytorch |
|
library_name: torchvision |
|
metrics: |
|
- accuracy |
|
--- |
|
# SuperBlock |
|
|
|
SuperBlock combines two techniques for efficient neural network training and inference: Supermask and Block Compressed Sparse Row (BSR) |
|
|
|
### Supermask |
|
[Supermask](https://arxiv.org/abs/2207.00670) is a technique for applying structured sparsity to neural networks using a learned mask. It works by learning a continuous mask (scores) that is applied element-wise to the weights of a neural network layer. The mask scores are learned separately from the weights and are thresholded based on a target sparsity level to obtain a binary mask. The mask determines which weigths are kept and which are pruned, and is learned during training. |
|
|
|
During inference, the binary mask is applied element-wise to the weights, pruning the weights that correspond to a 0 in the mask, resulting in a sparse network that can be efficiently computed. |
|
|
|
### Block compressed Sparse Row Format (BSR) |
|
[The BSR format](https://pytorch.org/docs/main/sparse.html#sparse-bsr-tensor) is a sparse matrix representation that stores dense sub-blocks of non-zero elements instead of individual non-zero elements. The matrix is divided into equal-sized blocks, and only the non-zero blocks are stored. |
|
|
|
The BSR format is efficient for sparse matrices with a block structure, where non-zero elements tend to cluster in dense sub-blocks. It reduces storage requirements and enables efficient matrix operations on the non-zero blocks. |
|
|
|
Currently, the BSR format is optimized for Nvidia A100 GPU(s) only. |
|
|
|
## Setup |
|
To use SuperBlock, you will need |
|
* [PyTorch](https://pytorch.org/get-started/locally/) |
|
|
|
To train the model or evaluate accuracy, you will need: |
|
* ImageNet2012-blurred dataset |
|
|
|
At least one GPU: |
|
* A100 or H100 |
|
|
|
## Installation |
|
* Clone this repo |
|
``` |
|
git clone https://github.com/pytorch-labs/superblock.git |
|
cd superblock |
|
``` |
|
* Create a new conda environment |
|
``` |
|
conda create -n superblock |
|
conda activate superblock |
|
``` |
|
* Install PyTorch. For best performance, we recommend `2.3.0.dev20240305+cu121` nightly |
|
``` |
|
pip install --pre torch==2.3.0.dev20240305+cu121 --index-url https://download.pytorch.org/whl/nightly/cu121 |
|
pip install --pre torchvision==0.18.0 --no-deps |
|
``` |
|
|
|
|
|
## Benchmarking |
|
Baseline: |
|
``` |
|
python benchmark.py \ |
|
--model vit_b_16 \ |
|
--batch-size 256 \ |
|
> /dev/null |
|
``` |
|
Result: |
|
``` |
|
532.1160546875 ms |
|
``` |
|
|
|
|
|
80% sparsity, block size 64 (random weights): |
|
``` |
|
python benchmark.py --model vit_b_16 \ |
|
--batch-size 256 \ |
|
--sparsity-linear 0.8 \ |
|
--sp-linear-tile-size 64 \ |
|
--sparsify-weights \ |
|
--bsr 64 \ |
|
> /dev/null |
|
``` |
|
Result: |
|
``` |
|
393.864453125 ms |
|
``` |
|
|
|
|
|
## Training |
|
Please refer to [TRAINING.md](TRAINING.md) for training from scratch. We use [Torchvision](https://github.com/pytorch/vision/tree/main/references/classification) as our framework for training. Supermask can be applied during training. |
|
|
|
To apply supermask, we have the following arguments at our disposal, |
|
|
|
* Apply Supermask to linear layers: |
|
``` |
|
--sparsity-linear |
|
--sp-linear-tile-size |
|
``` |
|
* Apply Supermask to conv1x1 layers: |
|
``` |
|
--sparsity-conv1x1 |
|
--sp-conv1x1-tile-size |
|
``` |
|
* Apply Supermask to all other convolutional layers: |
|
``` |
|
--sparsity-conv |
|
--sp-conv-tile-size |
|
``` |
|
* Skip the first transformer layer and/or last linear layer (ViT only): |
|
``` |
|
--skip-last-layer-sparsity |
|
--skip-first-transformer-sparsity |
|
``` |
|
|
|
For example, if you would like to train a `vit_b_16` from scratch using Supermask, you can use the respective torchvision command found in [TRAINING.md](TRAINING.md) and append the supermask arguments: |
|
``` |
|
torchrun --nproc_per_node=8 train.py\ |
|
--model vit_b_16 --epochs 300 --batch-size 512 --opt adamw --lr 0.003 --wd 0.3\ |
|
--lr-scheduler cosineannealinglr --lr-warmup-method linear --lr-warmup-epochs 30\ |
|
--lr-warmup-decay 0.033 --amp --label-smoothing 0.11 --mixup-alpha 0.2 --auto-augment ra\ |
|
--clip-grad-norm 1 --ra-sampler --cutmix-alpha 1.0 --model-ema\ |
|
--sparsity-linear 0.9 --sp-linear-tile-size 32 |
|
``` |
|
Through this command, we are training a `vit_b_16` with 90% sparsity to linear layers using 32x32 tiles. |
|
|
|
Please run `python train.py --help` for a full list of available arguments. |
|
|
|
## Evaluation |
|
|
|
To run an evaluation of a Supermask-trained model, you can use [evaluate.py](evaluate.py). Our current version has signficant speedup with float32 only and not float16, hence, to illustrate speedup, we don't pass `--amp` in the example commands below. |
|
|
|
``` |
|
MODEL_PATH=<put the path of the trained checkpoint here> |
|
IMAGENET_PATH=<put the path of ImageNet dataset here> |
|
NGPUS=1 # put number of available GPUS here |
|
``` |
|
|
|
* Offline sparsification with BSR: |
|
``` |
|
torchrun --nproc_per_node=${NGPUS} evaluate.py --model vit_b_16 --batch-size 256 --sparsity-linear 0.9 --sp-linear-tile-size 32 --weights-path ${MODEL_PATH} --data-path ${IMAGENET_PATH} --sparsify-weights --bsr 32 |
|
``` |
|
This command applies 90% sparsity to linear layers using 32x32 tiles, loads the model weights from ${MODEL_PATH}, loads the ImageNet validation set located at the specified path, applies offline sparsification to the weights, and converts the sparse weights to BSR format with a block size of 32. It is recommended to set `--bsr` the same as tile size. |
|
|
|
* Online sparsification without BSR: |
|
``` |
|
torchrun --nproc_per_node=${NGPUS} evaluate.py --model vit_b_16 --batch-size 256 --sparsity-linear 0.9 --sp-linear-tile-size 32 --weights-path ${MODEL_PATH} --data-path ${IMAGENET_PATH} |
|
``` |
|
This is similar to the previous command, but it does not apply offline sparsification or BSR conversion. Instead, the sparsity is applied on-the-fly during evaluation. |
|
|
|
Please run `python evaluate.py --help` for a full list of available arguments. |
|
|
|
Results (1x A100): |
|
* Baseline |
|
``` |
|
Test: Total time: 0:02:11 |
|
Test: Acc@1 78.392 Acc@5 93.592 |
|
``` |
|
|
|
* Sparsity= 0.9, Tile Size = 32, Online Sparsification, BSR = None |
|
``` |
|
Test: Total time: 0:01:52 |
|
Test: Acc@1 76.092 Acc@5 92.656 |
|
``` |
|
|
|
* Sparsity= 0.9, Tile Size = 32, Offline Sparsification, BSR = None |
|
``` |
|
Test: Total time: 0:01:54 |
|
Test: Acc@1 76.092 Acc@5 92.656 |
|
``` |
|
|
|
* Sparsity= 0.9, Tile Size = 32, Offline Sparsification, BSR = 32 |
|
``` |
|
Test: Total time: 0:01:25 |
|
Test: Acc@1 76.092 Acc@5 92.656 |
|
``` |
|
|
|
## Pretrained Weights |
|
|
|
### Download: |
|
Instead of training from scratch, if you'd like to use the Supermask weights of `vit_b_16` trained on privacy mitigated Imagenet-blurred, you can download them here: |
|
``` |
|
SPARSITY=0.80 # Checkpoints available for: 0.70, 0.80, 0.82, 0.84, 0.86, 0.88, 0.90 |
|
BLOCK_SIZE=32 # Checkpoints available for: 16, 32, 64 |
|
``` |
|
|
|
``` |
|
mkdir checkpoints |
|
# For baseline, |
|
wget https://huggingface.co/facebook/superblock-vit-b-16/resolve/main/checkpoints/baseline.pth -P checkpoints/ |
|
# For sparsified checkpoints, |
|
wget https://huggingface.co/facebook/superblock-vit-b-16/resolve/main/checkpoints/sp${SPARSITY}-ts${BLOCK_SIZE}.pth -P checkpoints/ |
|
``` |
|
|
|
### Benchmark: |
|
``` |
|
python benchmark.py --model vit_b_16 \ |
|
--batch-size 256 \ |
|
--sparsity-linear ${SPARSITY} \ |
|
--sp-linear-tile-size ${BLOCK_SIZE} \ |
|
--sparsify-weights \ |
|
--bsr ${BLOCK_SIZE} \ |
|
--weights-path ./checkpoints/superblock-vit-b-16-sp${SPARSITY}-ts${BLOCK_SIZE}.pth \ |
|
> /dev/null |
|
``` |
|
Result: |
|
``` |
|
530.342578125 ms |
|
``` |
|
|
|
### Evaluate: |
|
8 x A100 GPUs: |
|
``` |
|
torchrun --nproc_per_node=8 evaluate.py --model vit_b_16 --batch-size 256 --sparsity-linear ${SPARSITY} --sp-linear-tile-size ${BLOCK_SIZE} --bsr ${BLOCK_SIZE} --sparsify-weights --weights-path checkpoints/superblock-vit-b-16-sp${SPARSITY}-ts${BLOCK_SIZE}.pth --data-path ${IMAGENET_PATH} |
|
``` |
|
Result: |
|
``` |
|
Test: Total time: 0:01:01 |
|
Test: Acc@1 77.644 Acc@5 93.554 |
|
``` |
|
|
|
1 x A100 GPUs: |
|
``` |
|
torchrun --nproc_per_node=1 evaluate.py --model vit_b_16 --batch-size 256 --sparsity-linear ${SPARSITY} --sp-linear-tile-size ${BLOCK_SIZE} --bsr ${BLOCK_SIZE} --sparsify-weights --weights-path checkpoints/superblock-vit-b-16-sp${SPARSITY}-ts${BLOCK_SIZE}.pth --data-path ${IMAGENET_PATH} |
|
``` |
|
Result: |
|
``` |
|
Test: Total time: 0:01:51 |
|
Test: Acc@1 77.644 Acc@5 93.554 |
|
``` |
|
|
|
## License |
|
SuperBlock is released under the [MIT license](https://github.com/pytorch-labs/superblock?tab=MIT-1-ov-file#readme). |