logeshsenthil13/MobileNet_V1

MobileNetV1 Benchmark: PyTorch → ONNX → TensorRT (Jetson Orin Nano)

Every stage reports real ImageNet val top-1/top-5, parity vs PyTorch (same fixed 256 images), and latency. Artifacts land in outputs/mobilenet_v1/.

step1 PyTorch baseline ─▶ step2 ONNX export (FP32/FP16) ─▶ step3 ONNX vs PyTorch
        │
        └▶ make_calib ─▶ step4 ModelOpt INT8/INT4 ─▶ step3 (again, on quant ONNX)
                                   │
        step5 trtexec engines (ON JETSON) ─▶ step6 TRT vs PyTorch (ON JETSON)

Dev machine (RTX 6000 Ada)

source /data/users/logesh/QC_AI_HUB/venv/bin/activate
cd /data/users/logesh/QC_AI_HUB/MobileNet_V1/Proper

# 1. PyTorch baseline (accuracy + latency + reference logits)
python step1_pytorch.py --num-images 5000 --batch 32 --lat-batch 1
#   full val set: --num-images -1

# 2. ONNX export (FP32 + FP16)
python step2_export_onnx.py --fp16

# 3. ONNX Runtime validation (ONNX vs PyTorch)
python step3_eval_onnx.py --onnx outputs/mobilenet_v1/mobilenet_v1_fp32.onnx
python step3_eval_onnx.py --onnx outputs/mobilenet_v1/mobilenet_v1_fp16.onnx

# 4. Calibration + quantization (ModelOpt CLI)
python make_calib.py --num 512
bash step4_quantize.sh int8
bash step4_quantize.sh int4        # weight-only; see caveat below

# validate quantized ONNX (ONNX vs PyTorch)
python step3_eval_onnx.py --onnx outputs/mobilenet_v1/mobilenet_v1_int8.onnx
# NOTE: the INT4 ONNX cannot run in ONNX Runtime (ModelOpt emits TRT-style
# block quantization that ORT rejects) — validate INT4 at the TRT stage only.

Optional: local TRT check on the dev box

The dev box has trtexec 11.0 but python TensorRT 11.1 — engines built by that trtexec won't deserialize in python. Use the Python builder instead (Jetson doesn't have this problem; use step5 there):

python build_engine_py.py --onnx outputs/mobilenet_v1/mobilenet_v1_fp16.onnx \
    --engine outputs/mobilenet_v1/engines/mobilenet_v1_fp16.plan
python step6_eval_trt.py --engine outputs/mobilenet_v1/engines/mobilenet_v1_fp16.plan

Jetson Orin Nano

Copy over the ONNX files + reference artifacts + this folder's scripts (and mount/copy the val dataset, or eval a subset):

scp -r Proper/ jetson:~/mobilenet_v1_bench/
# needed on the Jetson: outputs/mobilenet_v1/{*.onnx,preprocess.json,ref_logits.npy,results_step1_pytorch.json}

On the Jetson (needs JetPack's tensorrt python, torch, numpy, pillow):

sudo nvpmodel -m 0 && sudo jetson_clocks       # stable clocks

# 5. Build engines (TRT engines are GPU-specific — must be built here)
bash step5_build_engines.sh fp32
bash step5_build_engines.sh fp16
bash step5_build_engines.sh int8
bash step5_build_engines.sh int4
# if trtexec not in PATH: TRTEXEC=/usr/src/tensorrt/bin/trtexec bash step5_build_engines.sh fp16

# 6. TRT accuracy + parity + latency (TRT vs PyTorch)
python3 step6_eval_trt.py --engine outputs/mobilenet_v1/engines/mobilenet_v1_fp32.plan
python3 step6_eval_trt.py --engine outputs/mobilenet_v1/engines/mobilenet_v1_fp16.plan
python3 step6_eval_trt.py --engine outputs/mobilenet_v1/engines/mobilenet_v1_int8.plan
python3 step6_eval_trt.py --engine outputs/mobilenet_v1/engines/mobilenet_v1_int4.plan

If the dataset lives elsewhere on the Jetson, edit DATASET_ROOT at the top of common.py.

Pure-TRT latency without Python overhead:

trtexec --loadEngine=outputs/mobilenet_v1/engines/mobilenet_v1_fp16.plan \
  --shapes=input:1x3x224x224 --warmUp=500 --duration=10 --useCudaGraph

Notes / caveats

• FP16 engine is built from the FP32 ONNX with --fp16 (standard TRT flow). The FP16 .onnx from step2 exists only to validate FP16 numerics in ONNX Runtime.
• INT8 uses ModelOpt explicit quantization: QDQ nodes, max calibration on 512 real val images, --high_precision_dtype fp32 (measured ~6pt better top-1 than the default fp16 conversion; ModelOpt auto-excludes the 13 depthwise convs from quantization). The engine is built with --int8 --fp16 on the Jetson so unquantized layers still run FP16. Measured INT8 top-1 drop vs FP32: ~2.3pt on a 512-image sample.
• step5_build_engines.sh auto-detects the trtexec generation: TRT 8.x/10.x (JetPack) uses --fp16/--int8 flags; TRT ≥ 11 is strongly typed (precision comes from the ONNX file), where the INT8 engine runs its unquantized layers in FP32 — expect the Jetson INT8-vs-FP16 speed gap to look better than the dev box's.
• INT4 caveat: ModelOpt/TensorRT INT4 is weight-only for MatMul/Gemm (transformer-style layers). MobileNetV1 is Conv-dominated, so INT4 touches at most the final classifier — expect ~FP16 performance and possibly a quantizer error/no-op. INT8 is the meaningful low-precision point for this CNN.
• The pretrained weights are timm's mobilenetv1_100.ra4_e3600_r224_in1k (modern training recipe, ~76% top-1 — well above the original paper's 70.9%). FP16 matches FP32 within ~0.1%.
• Parity metrics compare logits on a fixed seeded 256-image subset; latency is batch-1 by default (--lat-batch to change), accuracy on a seeded 5000-image subset by default (--num-images -1 for the full 50k).