aselimc Model Export & Optimization
PyTorch -> ONNX
torch.onnx.export(model, dummy_input, "model.onnx",
opset_version=17,
input_names=["input"], output_names=["output"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})
ONNX -> TensorRT
import tensorrt as trt
builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16) # or INT8
Validation
# Compare outputs
cos_sim = F.cosine_similarity(torch_out, onnx_out)
max_diff = (torch_out - onnx_out).abs().max()
assert cos_sim > 0.9999 and max_diff < 1e-3
Checklist
- Set model to
eval()mode before export - Use correct opset version (17+ recommended)
- Handle dynamic shapes (batch, sequence length)
- Validate numerical equivalence on 100+ samples
- Benchmark latency: PyTorch vs ONNX vs TensorRT
Key Libraries
torch.onnx, onnxruntime, tensorrt, torch.jit