This is the official PyTorch implementation of "Latency-aware Unified Dynamic Networks for Efficient Image Recognition", which is the extension of our NeurIPS 2022 paper: Latency-Aware Spatial-wise Dynamic Networks. The original LASNet code is at this URL.
We present Latency-aware Unified Dynamic Networks (LAUDNet), a unified framework that consolidates three representative dynamic paradigms: spatial-wise adaptive computation, dynamic layer skipping and dynamic channel skipping within a singular formulation. To accurately evaluate the practical latency of our model, we present a latency predictor that considers algorithms, scheduling strategies, hardware properties concurrently and accurately evaluates inference latency of dynamic operators. LAUDNet shows superior latency-accuracy tradeoff on a range of tasks (ImageNet classification, COCO object detection and instance segmentation) and a range of hardware devices (V100, RTX3090, RTX3060, TX2 and Nano).
This repo consists of three components: code for ImageNet classification, MMDetection detection & segmentation and latency predictor.
Main dependencies:
- Python: 3.9
- PyTorch: 1.13.1
- Torchvision: 0.14.1
- Timm: 0.6.12
See a sample training script for training details.
We implement the three dynamic-inference paradigms (i.e. token skipping, layer (block) skipping, and head (channel) skipping) based on the AdaViT repo.
Prerequisites:
- Prepare an ImageNet pretrained LAUDNet model.
- Setup a MMDetection-2.21.0 environment.
- Replace corresponding files in your mmcv environment with files in
mmcv_replace_file.
See a sample training script for training details.
Prerequisites:
- Prepare an ImageNet pretrained LAUDNet model.
- Setup a MMDetection-3.3.0 environment.
See a sample training script for training details.
See a sample evaluation script for evaluation details.
| model | Checkpoint Link |
|---|---|
| LAUD-ResNet101 channel-2222 target-0.5 | Tsinghua Cloud |
| LAUD-ResNet101 layer target-0.5 | Tsinghua Cloud |
@ARTICLE{han2024latency,
author={Han, Yizeng and Liu, Zeyu and Yuan, Zhihang and Pu, Yifan and Wang, Chaofei and Song, Shiji and Huang, Gao},
journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
title={Latency-aware Unified Dynamic Networks for Efficient Image Recognition},
year={2024},
volume={},
number={},
pages={1-17},
doi={10.1109/TPAMI.2024.3393530}
}If you have any questions, please feel free to contact the authors.
Yizeng Han: hanyz18@mails.tsinghua.edu.cn, yizeng38@gmail.com.
Zeyu Liu: liuzeyu20@mails.tsinghua.edu.cn, liuzeyu0020@gmail.com.
Zhihang Yuan: hahnyuan@gmail.com.
Yifan Pu: puyf23@mails.tsinghua.edu.cn.