Reproduce “TRT-ViT: TensorRT-oriented Vision Transformer” on OpenResearch

About this paper

Xin Xia, Jiashi Li, Jie Wu, Xing Wang, Xuefeng Xiao, Min Zheng, Rui Wang

Researchers from ByteDance Inc. developed TRT-ViT, a family of Vision Transformer architectures designed using NVIDIA TensorRT latency as a direct optimization metric for hardware-aware efficiency. These models achieve improved accuracy-latency trade-offs across image classification, segmentation, and object detection, with TRT-ViT-A outperforming ResNet50 by 2.2% Top-1 accuracy while exhibiting 10% less inference latency on a T4 GPU.

Key insights

• Placing Transformer blocks predominantly at later stages of the network, where feature map resolutions are lower, maximizes the trade-off between performance and hardware-measured efficiency.
• Adopting a 'shallow-then-deep' stage pattern, which increases the depth of later stages, can enhance model capacity and performance without significantly compromising inference efficiency.
• Hybrid blocks that mix Transformer and BottleNeck components are more efficient than pure Transformer blocks, with a 'global-then-local' internal arrangement (Transformer before BottleNeck) further improving performance.

Summary by alphaXiv.