leaplabthu / laudnet

IEEE TPAMI: Latency-aware Unified Dynamic Networks for Efficient Image Recognition

LAUDNet

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.

Introduction

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).

Usage

This repo consists of three components: code for ImageNet classification, MMDetection detection & segmentation and latency predictor.

ImageNet classification

CNNs

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.

Vision Transformers

We implement the three dynamic-inference paradigms (i.e. token skipping, layer (block) skipping, and head (channel) skipping) based on the AdaViT repo.

MMDetection detection & segmentation

RetinaNet, Faster-RCNN and MaskRCNN

Prerequisites:

1. Prepare an ImageNet pretrained LAUDNet model.
2. Setup a MMDetection-2.21.0 environment.
3. Replace corresponding files in your mmcv environment with files in mmcv_replace_file.

See a sample training script for training details.

DDQ-DETR and Mask2Former

Prerequisites:

1. Prepare an ImageNet pretrained LAUDNet model.
2. Setup a MMDetection-3.3.0 environment.

See a sample training script for training details.

Latency predictor

See a sample evaluation script for evaluation details.

Performance

Model Zoo

model	Checkpoint Link
LAUD-ResNet101 channel-2222 target-0.5	Tsinghua Cloud
LAUD-ResNet101 layer target-0.5	Tsinghua Cloud

Citation

@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}
}

Contact

If you have any questions, please feel free to contact the authors.

Yizeng Han: [email protected], [email protected].

Zeyu Liu: [email protected], [email protected].

Zhihang Yuan: [email protected].

Yifan Pu: [email protected].

laudnet's Issues

如何实现加速？

作者您好。从库里面的代码来看，对通道、空间维度的动态裁剪都是通过卷积之后apply mask实现，这种方式相对于baseline增加了计算量，在实际的运行中时怎样实现速度提升的呐？？谢谢。