README.md

Export 🚛

Stage 1 of the Model Gym pipeline — converts an imported PyTorch model into a deployable format.

ImportedModel (nn.Module)
    │
    ├── OnnxExporter      → model.onnx   (any runtime)
    │
    └── TensorRTBuilder   → model.engine (NVIDIA GPU only)

Modules

Module	Class	Input	Output
onnx_exporter.py	OnnxExporter	nn.Module	.onnx
tensorrt_builder.py	TensorRTBuilder	.onnx	.engine

OnnxExporter

Wraps torch.onnx.export with dynamic axes, shape inference, and model validation.

Features

• Dynamic batch and sequence axes — one export works for any batch size at runtime
• Opset pinning (default 17 — widest TensorRT support)
• Automatic onnx.shape_inference.infer_shapes pass after export
• onnx.checker.check_model validation — fails fast on malformed graphs
• HuggingFace convenience method — load + export in one call

CLI

pip install -e ".[export]"

# From a HuggingFace model ID
python -m src.export.onnx_exporter \
    --model  bert-base-uncased \
    --task   text-classification \
    --output models/bert.onnx \
    --opset  17 \
    --seq-len 128

# Verbose torch.onnx output
python -m src.export.onnx_exporter \
    --model bert-base-uncased --output models/bert.onnx --verbose

Python API

from src.export.onnx_exporter import OnnxExporter

# Option A — convenience method (HuggingFace model)
OnnxExporter.from_huggingface(
    model_id="bert-base-uncased",
    task="text-classification",
    output_path="models/bert.onnx",
    opset=17,
    seq_len=128,
)

# Option B — bring your own model
import torch
from src.export.onnx_exporter import OnnxExporter

exporter = OnnxExporter(opset=17, dynamic_batch=True, dynamic_seq=True)
exporter.export(
    model=my_model,
    output_path="models/custom.onnx",
    input_shape=(1, 128),  # (batch, seq_len) for synthetic dummy input
)

# Option C — custom dummy inputs (e.g. multimodal model)
exporter.export(
    model=my_model,
    output_path="models/custom.onnx",
    dummy_inputs={
        "input_ids": torch.zeros((1, 128), dtype=torch.int64),
        "attention_mask": torch.ones((1, 128), dtype=torch.int64),
        "token_type_ids": torch.zeros((1, 128), dtype=torch.int64),
    },
)

Options

Argument	Default	Description
--model	required	HuggingFace model ID or local path
--task	text-classification	HuggingFace task string
--output	models/model.onnx	Destination .onnx file
--opset	17	ONNX opset version
--seq-len	128	Dummy input sequence length
--verbose	off	Pass verbose to torch.onnx.export

---

TensorRTBuilder

Builds a serialised .engine file from an ONNX model using the TensorRT Python API.

Requires: NVIDIA GPU + TensorRT + CUDA toolkit.

Features

• FP32 / FP16 / INT8 precision flags
• Dynamic shape profiles — min / opt / max (batch, seq_len) for kernel selection
• INT8 calibration cache input via _Int8CacheCalibrator
• Builder workspace cap (default 4 GB)
• Graceful ImportError with install instructions when TensorRT is absent

CLI

# FP16 engine — most common for production
python -m src.export.tensorrt_builder \
    --input      models/bert.onnx \
    --output     models/bert_fp16.engine \
    --precision  fp16 \
    --workspace-gb 4

# INT8 engine with calibration cache
python -m src.export.tensorrt_builder \
    --input      models/bert.onnx \
    --output     models/bert_int8.engine \
    --precision  int8 \
    --calib-table data/calib_bert.cache

# Custom shape profile (batch 1–16, seq 128)
python -m src.export.tensorrt_builder \
    --input     models/bert.onnx \
    --output    models/bert.engine \
    --min-batch 1 --opt-batch 8 --max-batch 16 \
    --seq-len   128

Python API

from src.export.tensorrt_builder import TensorRTBuilder, ShapeProfile

# FP16 engine with custom shape profile
builder = TensorRTBuilder(
    precision="fp16",
    workspace_gb=4.0,
    profiles=[
        ShapeProfile(
            input_name="input_ids",
            min=(1, 32),
            opt=(8, 128),
            max=(32, 128),
        )
    ],
)
builder.build("models/bert.onnx", "models/bert_fp16.engine")

# INT8 with calibration cache
builder = TensorRTBuilder(
    precision="int8",
    calib_table="data/calib_bert.cache",
)
builder.build("models/bert.onnx", "models/bert_int8.engine")

Options

Argument	Default	Description
--input	required	Source .onnx file
--output	required	Destination .engine file
--precision	fp16	fp32 / fp16 / int8
--workspace-gb	4	Builder workspace limit in GB
--calib-table	—	INT8 calibration cache (required for int8)
--min-batch	1	Min batch size for shape profile
--opt-batch	8	Opt batch size for shape profile
--max-batch	32	Max batch size for shape profile
--seq-len	128	Sequence length for shape profile
--verbose	off	TensorRT verbose logging

---

End-to-end example

# 1. Import
python -m src.importer.huggingface \
    --model bert-base-uncased --task text-classification

# 2. Export to ONNX
python -m src.export.onnx_exporter \
    --model bert-base-uncased --output models/bert.onnx

# 3. Build TensorRT FP16 engine  (GPU required)
python -m src.export.tensorrt_builder \
    --input models/bert.onnx \
    --output models/bert_fp16.engine \
    --precision fp16

# 4. Benchmark
python -m src.benchmark.harness \
    --backends onnx trt \
    --model models/bert.onnx \
    --batch-sizes 1 8 32 \
    --out results/bench_bert.json

Dependencies

# pyproject.toml — [project.optional-dependencies]
export = ["torch", "onnx", "onnxruntime", "tensorrt"]

pip install -e ".[export]"   # torch + onnx always; tensorrt only on GPU hosts

tensorrt import is guarded — the package installs and imports cleanly on CPU-only machines; the ImportError is only raised when TensorRTBuilder.build() is called.

File Layout

src/export/
├── __init__.py           exports OnnxExporter, TensorRTBuilder
├── onnx_exporter.py      OnnxExporter class + CLI
├── tensorrt_builder.py   TensorRTBuilder + ShapeProfile + CLI
└── README.md             this file

---

© 2026 Hitesh Kumar Sahu · Licensed under Apache 2.0