Skip to content

Tutorial - EfficientViT

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

What you need

  1. One of the following Jetson devices:

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

  2. Running one of the following versions of JetPack:

    JetPack 5 (L4T r35.x)

  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.