Vision Transformer Implementation on Edge GPU (AGX Orin) for Image Classification

With the use of vision transformers, the transformer architecture, which was first designed for natural language processing (NLP), has been adapted for computer vision applications. This paper introduces a robust approach for flower image classification leveraging a pre-trained Vision Transformer (ViT) model. The approach is optimized for flower classification tasks by using the Vision Transformer as the base model. Training images were enlarged and resized to improve the model's generalization skills. Furthermore, during the training phase, particular training parameters were set up to thoroughly assess the model's categorization accuracy. So, in this model we are basically classifying images into two classes of flowers that are daisy and dandelion. We have implemented this model on the NVIDIA AGX ORIN 64 GB developer kit. There are three power modes available in this kit i.e. 15W, 30W and 50 W. We have operated our GPU kit in all three modes and analyzed for all the parameters including power, memory utilization, CPU core utilization and accuracy etc. We have achieved 98.58 % training accuracy and 97.32 % testing accuracy while training our pre-trained model.