FKD: A Fast Knowledge Distillation Framework for Visual Recognition
Official PyTorch implementation of paper A Fast Knowledge Distillation Framework for Visual Recognition. Zhiqiang Shen and Eric Xing from CMU and MUZUAI.
Abstract
Knowledge Distillation (KD) has been recognized as a useful tool in many visual tasks, such as the supervised classification and self-supervised representation learning, while the main drawback of a vanilla KD framework lies in its mechanism that most of the computational overhead is consumed on forwarding through the giant teacher networks, which makes the whole learning procedure in a low-efficient and costly manner. In this work, we propose a Fast Knowledge Distillation (FKD) framework that simulates the distillation training phase and generates soft labels following the multi-crop KD procedure, meanwhile enjoying the faster training speed than ReLabel as we have no post-processes like RoI align and softmax operations. Our FKD is even more efficient than the conventional classification framework when employing multi-crop in the same image for data loading. We achieve 79.8% using ResNet-50 on ImageNet-1K, outperforming ReLabel by ~1.0% while being faster. We also demonstrate the efficiency advantage of FKD on the self-supervised learning task.
Supervised Training
Preparation
-
Install PyTorch and ImageNet dataset following the official PyTorch ImageNet training code. This repo has minimal modifications on that code.
-
Download our soft label. We provide multiple types of soft labels, and we recommend to use Marginal Smoothing Top-5 (500-crop).
FKD Training on CNNs
To train a model, run train_FKD.py with the desired model architecture and the path to the soft label and ImageNet dataset:
python train_FKD.py -a resnet50 --lr 0.1 --num_crops 4 -b 1024 --cos --softlabel_path [soft label path] [imagenet-folder with train and val folders]
For --softlabel_path, simply use format as ./FKD_soft_label_500_crops_marginal_smoothing_k_5
Multi-processing distributed training is supported, please refer to official PyTorch ImageNet training code for details.
Evaluation
python train_FKD.py -a resnet50 -e --resume [model path] [imagenet-folder with train and val folders]
Trained Models
| Model | accuracy (Top-1) | weights | configurations |
|---|---|---|---|
ReLabel ResNet-50 |
78.9 | -- | -- |
FKD ResNet-50 |
79.8 | link | Table 10 in paper |
ReLabel ResNet-101 |
80.7 | -- | -- |
FKD ResNet-101 |
81.7 | link | Table 10 in paper |
FKD Training on ViT/DeiT and SReT
To train a ViT model, run train_ViT_FKD.py with the desired model architecture and the path to the soft label and ImageNet dataset:
cd train_ViT
python train_ViT_FKD.py -a SReT_LT --lr 0.002 --wd 0.05 --num_crops 4 -b 1024 --cos --softlabel_path [soft label path] [imagenet-folder with train and val folders]
For the instructions of SReT_LT model, please refer to SReT for details.
Evaluation
python train_ViT_FKD.py -a SReT_LT -e --resume [model path] [imagenet-folder with train and val folders]
Trained Models
| Model | FLOPs | #params | accuracy (Top-1) | weights | configurations |
|---|---|---|---|---|---|
DeiT-T-distill |
1.3B | 5.7M | 74.5 | -- | -- |
FKD ViT/DeiT-T |
1.3B | 5.7M | 75.2 | link | Table 11 in paper |
SReT-LT-distill |
1.2B | 5.0M | 77.7 | -- | -- |
FKD SReT-LT |
1.2B | 5.0M | 78.7 | link | Table 11 in paper |
Fast MEAL V2
Please see MEAL V2 for the instructions to run FKD with MEAL V2.
Self-supervised Representation Learning Using FKD
Please see FKD-SSL for the instructions to run FKD code for SSL task.
Citation
@article{shen2021afast,
title={A Fast Knowledge Distillation Framework for Visual Recognition},
author={Zhiqiang Shen and Eric Xing},
year={2021},
journal={arXiv preprint arXiv:2112.01528}
}
Contact
Zhiqiang Shen (zhiqians at andrew.cmu.edu or zhiqiangshen0214 at gmail.com)
47 Jan 9, 2023
9 Nov 14, 2022
2 Oct 20, 2022
40 Dec 22, 2022
112 Dec 2, 2022
19 Jun 7, 2022
9 Sep 26, 2022
180 Dec 19, 2022
11 Nov 12, 2022
11 Nov 22, 2022
147 Dec 3, 2022
75 Dec 28, 2022
80 Dec 16, 2022
215 Dec 19, 2022
26 Oct 13, 2022
13 Dec 10, 2022
47 Nov 17, 2022
9 Dec 12, 2022