A ViT-based lightweight method for the UAV platform object detection tasks
Zhi Fang, Zhizhong Xi, Mengen Xu, Xihui Fan · 4th International Conference on Information Science, Electrical, and Automation Engineering (ISEAE 2022) · 2022
Lightweight design is a key way to realize the engineering application of deep learning algorithms on Unmanned Aerial Vehicle (UAV) platform. Aiming at the low detection accuracy of the current real-time object detection algorithms of the UAV platform, a lightweight model based on Vision Transformer (ViT) is designed in this paper. Firstly, a small Convolutional Neural Network (CNN) is used to extract primary features for reducing the number of parameters and computation amount of ViT network in this model, and using window modeling to replace part of the global modeling. Then, a feature-level mask self-supervised training method is applied to pre-train the ViT structure, which helps to accelerate the convergence and avoid a lot of labeling work. Finally, the result compared with other UAV lightweight object detection algorithms in the visdrone2018 dataset shows that this model has higher average accuracy on ensuring real-time speed, and verifies the effectiveness and reference value of the lightweight design method proposed in this paper.