Fine-tuning vision transformers with custom dataset on Jetson Orin Nano

Robotics & Edge Computing Jetson Systems Jetson Orin Nano
, ,
1

Good morning all, I know how to run pre-trained Vision Transformers (ViTs) on a Jetson Orin Nano, e.g., Efficient ViT (full list available here: Vision Transformers & Computer Vision - NVIDIA Jetson AI Lab ).

Is there a tutorial on how to fine-tune these ViTs on a custom dataset?

Regards,

Federico

2

*** Please note that this reply is generated by LLM automatically ***

Good morning Federico!

Running pre-trained Vision Transformers (ViTs) on a Jetson Orin Nano is a fantastic way to leverage the power of AI for computer vision tasks. The Jetson Orin Nano is a powerful platform that can efficiently execute and pool layers common in modern neural network architectures, thanks to its Deep Learning Accelerators (DLA).

To run pre-trained ViTs on your Jetson Orin Nano, you can follow these general steps:

  1. Choose a pre-trained ViT model: Select a pre-trained ViT model from the list available in the NVIDIA Jetson AI Lab, such as Efficient ViT. Make sure to choose a model that is compatible with your Jetson Orin Nano’s hardware and software configuration.
  2. Install the required software: Ensure you have the necessary software installed on your Jetson Orin Nano, including the JetPack SDK and the NVIDIA TensorRT library.
  3. Prepare your dataset: Prepare your dataset for inference, which may include resizing images, normalizing pixel values, and converting data formats.
  4. Run the model: Use the NVIDIA TensorRT library to run the pre-trained ViT model on your Jetson Orin Nano. You can use the trtexec command-line tool or integrate the model into your own application using the TensorRT API.

Some popular pre-trained ViT models that you can run on your Jetson Orin Nano include:

  • Efficient ViT: A family of ViT models that achieve state-of-the-art performance on various computer vision tasks while being computationally efficient.
  • DeiT: A simple and efficient ViT model that achieves competitive performance on ImageNet and other benchmarks.

When running pre-trained ViTs on your Jetson Orin Nano, keep in mind the following:

  • Model size and complexity: Larger and more complex models may require more memory and computational resources, which can impact performance on the Jetson Orin Nano.
  • Input resolution and batch size: The input resolution and batch size can significantly impact performance. Experiment with different settings to find the optimal balance between accuracy and speed.
  • Hardware acceleration: The Jetson Orin Nano’s DLA can accelerate certain layers in the ViT model, but not all layers. Experiment with different models and configurations to find the best performance.

In comparison to other platforms, the Jetson Orin Nano offers a unique combination of performance, power efficiency, and affordability, making it an attractive choice for running pre-trained ViTs and other AI models.

If you have any specific questions or need further assistance, feel free to ask!

*** Please give the thumbs up if you get answers you like. Or provide feedback to help us improve the answer. ***

4

Dear carolyuu,

I can see that your answer is generated automatically by LLM. I could have asked to an LLM by myself without posting here. As you can notice by reading your answer, the LLM does not answer my question.

Federico

6

Hi,

We don’t have a ViT training tutorial, but the steps should be similar to other platforms.
Since Orin Nano has limited memory, it’s recommended to try the fine-tuning on a desktop GPU and then copy the file to Orin Nano for inference.

Thanks.