EdgeTAM PTQ

English | 中文

ONNX export and post-training quantization (PTQ) toolchain based on EdgeTAM, supporting ONNX model export and quantized deployment via the Horizon HBDK toolchain.

Project: https://github.com/D-Robotics-AI-Lab/EdgeTAM_PTQ

Installation

git clone https://github.com/D-Robotics-AI-Lab/EdgeTAM_PTQ.git && cd EdgeTAM_PTQ
pip install -e .
pip install onnx onnxruntime matplotlib pillow

The Horizon HBDK/HBRuntime toolchain (horizon_tc_ui) is required for quantized inference.

Model Preparation

Download the pretrained checkpoint from the EdgeTAM repository and place it at ./checkpoints/edgetam.pt.

If checkpoint key names do not match the code, use convert_weights.py:

python convert_weights.py --src ./checkpoints/edgetam.pt
# Produces ./checkpoints/edgetam_converted.pt

Pipeline Overview

① export_onnx_box_mask.py               → Export ONNX models
        ↓
② test_onnx_video_memory_box.py         → Float prefill, generate .npy conditioning memory
        ↓
③ test_onnx_video_track.py              → Float ONNX Runtime video tracking validation
        ↓
④ Quantize model_prompt_to_memory_points → PTQ quantize prefill model (produce .bc)
⑤ Quantize model_track_step_s7           → PTQ quantize tracking model (produce .bc)
        ↓
⑥ test_quantionnx_video_prompt.py   → Quantized prefill accuracy validation
⑥ test_quantionnx_video_track_nv12.py   → Quantized tracking NV12 inference validation

Workflow

The pipeline consists of six steps, executed in order:

① Export ONNX Models

Export EdgeTAM to two ONNX model files:

python3 export_onnx_box_mask.py

Outputs:

File	Description
model_prompt_to_memory_points.onnx	Conditioning frame (frame 0) prompt→memory bank model. Inputs: image + point_coords + point_labels. Outputs: low_res_masks_logits + conditioning memory (maskmem_features / maskmem_pos_enc / obj_ptr)
model_track_step_s7.onnx	Tracking frame model with memory slots=7. Outputs: updated mask + memory for per-frame tracking from frame 1 onward

HBDK compatibility post-processing is applied automatically: IR version downgrade to 9 and ScatterND reduction attribute removal.

② Float ONNX Inference — Generate Conditioning Frame Reference

Run prefill with ONNX Runtime and save output tensors as initial memory for subsequent tracking:

python3 test_onnx_video_memory_box.py --mode points   # point prompt
python3 test_onnx_video_memory_box.py --mode box      # box prompt

Outputs:

File	Description
float_memory_low_res_logits0.npy	Frame 0 low-resolution mask logits
float_memory_cond_maskmem_features.npy	Conditioning memory features (slot 0)
float_memory_cond_maskmem_pos_enc.npy	Conditioning memory positional encoding (slot 0)
float_memory_cond_obj_ptr.npy	Conditioning object pointer (slot 0)
result_float.png	Frame 0 visualization

③ Float ONNX Inference — Video Tracking

Load the .npy conditioning memory from step ② and perform per-frame tracking:

python3 test_onnx_video_track.py --mode points --num_max_frames 50

Arguments:

• --mode: Prompt type (points / box / mask)
• --video_dir: Directory of video frames (default ./notebooks/videos/bedroom, sorted by filename e.g. 00000.jpg, 00001.jpg)
• --num_max_frames: Maximum number of tracking frames
• --out_dir: Output directory for results

④ Quantize — Conditioning Frame Prefill

RDK S100 Quantization

Quantization on RDK S100 GPU docker toolchain (OE>=3.7.0).

hb_compile --config quantification/model_prompt_to_memory_points_config.yaml

RDK S600 Quantization

Quantization on RDK S600 GPU docker toolchain (OE>=3.7.0).

hb_compile --config quantification/model_prompt_to_memory_points_s600_config.yaml

⑤ Quantize — Tracking Frame

Quantization on RDK S100 GPU docker toolchain (OE>=3.7.0).

# env dependence
dpkg -i quantification/hmct_gpu-2.7.2.post0.dev202603261636+da214c21c528d840df6aca55c715877baba03ea2.cu128-cp310-cp310-linux_x86_64.whl

hb_compile --config quantification/model_prompt_to_memory_points_config.yaml

RDK S600 量化

Quantization on RDK S600 GPU docker toolchain (OE>=3.7.0).

# env dependence
dpkg -i quantification/hmct_gpu-2.7.2.post0.dev202603261636+da214c21c528d840df6aca55c715877baba03ea2.cu128-cp310-cp310-linux_x86_64.whl

hb_compile --config quantification/model_prompt_to_memory_points_s600_config.yaml

⑥ Quantized Inference

Use the .bc models from steps ④⑤ to validate quantized inference accuracy.

Quantized prefill test:

python3 test_quantionnx_video_prompt_nv12.py --mode points

This script runs both float ONNX Runtime and HBRuntime quantized inference, comparing cosine similarity to measure quantization accuracy loss.

Quantized tracking test (NV12 input):

python3 test_quantionnx_video_track_nv12.py \
    --mode points \
    --num_max_frames 50 \
    --quanti_model ./model_output/model_track_step_s7_quantized_model.bc

Arguments:

• --quanti_model: Path to the quantized .bc model file
• --low_res_logits0_npy and similar: Paths to .npy files from step ② for loading conditioning memory
• Image input is automatically converted to NV12 format (Y + interleaved UV) to match Horizon hardware encoding requirements

License

Based on EdgeTAM code, licensed under Apache 2.0.