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:

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.

Prerequisites:

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

See a sample training script for training details.

See a sample evaluation script for evaluation details.

@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: [email protected], [email protected].

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

Zhihang Yuan: [email protected].

Yifan Pu: [email protected].