Skip to content

Tutorial - EfficientViT

Let's run MIT Han Lab's EfficientViT on Jetson!

What you need

  1. One of the following Jetson:

    Jetson AGX Orin 64GB Jetson AGX Orin (32GB) Jetson Orin Nano Orin (8GB)

  2. Running one of the following JetPack.5x

    JetPack 5.1.2 (L4T r35.4.1) JetPack 5.1.1 (L4T r35.3.1) JetPack 5.1 (L4T r35.2.1)

  3. Sufficient storage space (preferably with NVMe SSD).

    • 10.9 GB for efficientvit container image
    • Space for checkpoints

Clone and set up jetson-containers

git clone https://github.com/dusty-nv/jetson-containers
cd jetson-containers
sudo apt update; sudo apt install -y python3-pip
pip3 install -r requirements.txt

How to start

Use run.sh and autotag script to automatically pull or build a compatible container image.

cd jetson-containers
./run.sh $(./autotag efficientvit)

Usage of EfficientViT

The official EfficientViT repo shows the complete usage information.

https://github.com/mit-han-lab/efficientvit#usage

Run example/benchmark

Inside the container, a small benchmark script benchmark.py is added under /opt/efficientvit directory by the jetson-container build process.

It is to test EfficientViT-L2-SAM in bounding box mode, so we can use this as an example and verify the output.

Download l2.pt model

mkdir -p /data/models/efficientvit/sam/
cd /data/models/efficientvit/sam/
wget https://huggingface.co/han-cai/efficientvit-sam/resolve/main/l2.pt

The downloaded checkpoint file is stored on the /data/ directory that is mounted from the Docker host.

Run benchmark script

cd /opt/efficientvit
python3 ./benchmark.py

At the end you should see a summary like the following.

AVERAGE of 2 runs:
  encoder --- 0.062 sec
  latency --- 0.083 sec
Memory consumption :  3419.68 MB

Check the output/result

The output image file (of the last inference result) is stored as /data/benchmarks/efficientvit_sam_demo.png.

It is stored under /data/ directory that is mounted from the Docker host.
So you can go back to your host machine, and check ./jetson-containers/data/benchmark/ directory.

You should find the output like this.