Adaptive class token knowledge distillation for efficient vision transformer

Official PyTorch implementation of the paper:

"Adaptive Class Token Knowledge Distillation for Efficient Vision Transformer"
Published in Knowledge-Based Systems (Elsevier, 2024)
📄 DOI: 10.1016/j.knosys.2024.112531
🔗 Journal link: ScienceDirect

Introduction

We propose [CLS]-KD, an adaptive knowledge distillation method for vision transformers that leverages class tokens and CLS-patch attention maps to enhance performance.

Usage

Requirements

pytorch==1.8.0
timm==0.5.4

Data preparation

Download and extract ImageNet train and val images from http://image-net.org/. The directory structure is:

│path/to/imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

Training on ImageNet-1K

To train a DeiT-Tiny student with a DeiT-Base teacher, run:

python -m torch.distributed.launch --nproc_per_node=4 main.py --distributed --output_dir <output-dir> --data-path <dataset-dir> --teacher-path <path-of-teacher-checkpoint> --model deit_tiny_patch16_224 --teacher-model deit_base_patch16_224 --distillation-type soft --distillation-alpha 1 --distillation-beta 1 --manifold --w-cls 1.0 --w-atn 2.0 --last-w 4.0 --K 192 --s-id 0 1 2 3 8 9 10 11 --t-id 0 1 2 3 8 9 10 11 --drop-path 0 --batch-size 128 --seed 0

Note: pretrained deit_base_patch16_224 model can be download from deit.

Test on ImageNet-1K

To test a DeiT-Tiny student with a DeiT-Base teacher, run:

python main.py --eval --output_dir <output-dir> --data-path <dataset-dir> --teacher-path <path-of-teacher-checkpoint> --model deit_tiny_patch16_224 --teacher-model deit_base_patch16_224 --distillation-type soft --distillation-alpha 1 --distillation-beta 1 --manifold --w-cls 1.0 --w-atn 2.0 --last-w 4.0 --K 192 --s-id 0 1 2 3 8 9 10 11 --t-id 0 1 2 3 8 9 10 11 --drop-path 0 --resume <path-of-distilled-model>

Result

Teacher	Student	Acc@1	Checkpoint & log
DeiT-Base	DeiT-Tiny	75.1	checkpoint / log
DeiT-Base	DeiT-Small	81.49	checkpoint / log

Citation

If you find this project useful in your research, please consider cite:

@article{Kang2024clskd,
title = {Adaptive class token knowledge distillation for efficient vision transformer},
author = {Minchan Kang and Sanghyeok Son and Daeshik Kim},
journal = {Knowledge-Based Systems},
volume = {304},
pages = {112531},
year = {2024},
publisher={Elsevier}
}

Acknowledgment

This repo is based on DeiT, manifold-distillation and pytorch-image-models.

Contact

Email : mc.kang@kaist.ac.kr