Lab 02 — Port PyTorch → ONNX → TensorRT

Goal: take a PyTorch model, export it to ONNX, compile it to a TensorRT engine, validate numerical correctness against the FP16 reference, and keep a log of every op/shape that broke and how you fixed it. That log is the deliverable — it's the JD's "model conversion workflows" and "issue triage" made tangible.

Knowledge prereq: K03 — Model Conversion & Compilers.

Stack: torch, onnx, onnxruntime, onnxsim, tensorrt + trtexec (NVIDIA), Netron (graph viewer). The ONNX half runs anywhere; TensorRT needs NVIDIA.

Steps

  1. Pick a model with some teeth — not a toy MLP. A vision model (ResNet/YOLO) is the gentle path; a small transformer (BERT, or a tiny Llama) exposes the dynamic-shape pain that's the real lesson.

  2. Export to ONNX: torch.onnx.export(..., opset_version=17, dynamic_axes={...}) or torch.onnx.dynamo_export. Record the opset and which axes you marked dynamic.

  3. Inspect in Netron: open the .onnx, eyeball the graph. Run onnx.checker.check_model and onnxsim (simplify/constant-fold). Note anything weird (extra Cast/Transpose nodes, unfused attention).

  4. Validate the ONNX against PyTorch: run both on the same inputs via ONNX Runtime; assert np.allclose(rtol=1e-3, atol=1e-3). Fix any divergence here before TensorRT — the bug is usually in the export, not the compiler (K03 §10).

  5. Compile to TensorRT: trtexec --onnx=model.onnx --saveEngine=model.plan --fp16 (add --int8 + a calibration cache for the quantization variant). For dynamic shapes, pass --minShapes/--optShapes/--maxShapes (optimization profiles).

  6. Hit the breakages (K03 §6) and log each one:

    SymptomClass (op / shape / numeric)Root causeFix
    e.g. "unsupported op ScatterND"opcustom indexingrewrite with gather/slice
    e.g. "engine fails on batch>1"shapestatic exportadd optimization profile
    e.g. "logits differ by 0.3"numericFP16 overflow in a layerkeep that layer FP32

  7. Validate the engine: compare TensorRT outputs vs the FP16 PyTorch reference (max abs error, cosine of logits). If it diverges, bisect: FP32 engine first (isolates a real bug from FP16 drift), then layer-by-layer.

  8. Measure the win: latency/throughput of PyTorch-eager vs torch.compile vs TensorRT. Confirm it's actually faster (sometimes a fallback-to-CPU op makes it slower — measure, don't assume).

  9. Write the port report (K03 §10): ops fixed, max error, accuracy delta, speedup — the customer-ready artifact.

What you must be able to explain

  • Why an engine built on one GPU SKU may not load on another (K03 §4).
  • Why dynamic shapes are the hard part and how optimization profiles solve them.
  • The difference between "it runs" and "it's correct" — and how you proved correct.

Result / résumé bullet

"Ported a transformer from PyTorch to TensorRT via ONNX: resolved 3 unsupported-op and dynamic-shape failures, validated numerical parity (cosine 0.9998 vs FP16 reference), and delivered a 2.4× latency reduction with a documented conversion-issue log and port report."