Skip to main content

TensorRT backend for embedl-deploy.

Project description

embedl-deploy

Python package to make AI models deployment-ready for any hardware.

Why embedl-deploy

PyTorch models are flexible, but edge hardware is not. Hardware toolchains may fail due to unsupported operators, apply implicit transformations and fusions during compilation and quantization leading to deployment issues.

embedl-deploy eliminates these surprises by enforcing hardware and compiler constraints directly in PyTorch, so what you build, train, and debug is what actually runs on the device. It converts your models to be compatible for the hardware target ensuring correct quantization and compilation.

Features

  • Hardware-accurate PyTorch Intermediate Representation (IR): Build models using a hardware-aware PyTorch intermediate representation that mirrors the behavior of the compiled artifact, e.g., fused convolutions. Unsupported operators and compatibility issues are surfaced early and resolved explicitly before compilation, within PyTorch.

  • Quantization: Supports post-training quantization (PTQ) and quantization-aware-training (QAT). Fake quantization is applied in PyTorch with explicit quantization operator placements. PTQ methods are included and QAT can be applied directly to the transformed and quantized models with no additional dependencies required.

  • Guaranteed deployable artifacts: Produce optimized compilation artifacts ready for deployment on the target device with predictable performance and accuracy.

Supported Backends

Backend Status
NVIDIA TensorRT (v10.3) Supported
Lattice SensAI (v8.0) In Development

Contact Embedl for other backends.

Installation

pip install "embedl-deploy[tensorrt]"

Note that you may need to also install onnx and onnx-simplifier to export and get the exported model compiled with TensorRT if using ONNX as an intermediate.


Quick Start for TensorRT Backend

import torch
from embedl_deploy import transform
from embedl_deploy.quantize import quantize
from embedl_deploy.tensorrt import TENSORRT_PATTERNS
from torchvision.models import resnet18 as Model

# 1. Load a standard PyTorch model
model = Model().eval()
example_input = torch.randn(1, 3, 224, 224)

# 2. Transform — fuse and optimize for TensorRT in one call
# For more compatibility you can trace your model with torch.export.export
# as follows:
# model = torch.export.export(model, (example_input)).module()
res = transform(model, patterns=TENSORRT_PATTERNS)
print("Model\n", res.model.print_readable())
print("Matches", "\n".join([str(match) for match in res.matches]))


# 3. Quantize (PTQ)
def calibration_loop(model: torch.fx.GraphModule):
    model.eval()
    for _ in range(100):
        model(torch.randn(1, 3, 224, 224))


quantized_model = quantize(
    res.model, (example_input,), forward_loop=calibration_loop
)
quantized_model.eval()

# 4. Export as usual (dynamo exported models may have compilation issues)
torch.onnx.export(
    quantized_model, (example_input,), "model.onnx", dynamo=False
)

# 5. Quantization-aware training with a training loop
qat_model = quantized_model.train()
# Freeze BatchNorm, or apply other QAT utilities as needed
# train(qat_model)

Compile

Compilation can be done with TensorRT's trtexec tool, which can take the ONNX model and compile it for inference. The exported layer info and profile can be used for debugging, optimization and visualization.

Note: that the ONNX model might need to be simplified with onnx-simplifier to make trtexec compile it. Dynamo exported models may have compilation issues, so it's recommended to export with dynamo=False.

onnxsim model.onnx model.onnx
/usr/src/tensorrt/bin/trtexec --onnx=model.onnx --fp16 --int8 --useCudaGraph

Optionally you can get the layer profile with the following flags:

--exportLayerInfo=layer_info.json
--exportProfile=profile.json
--profilingVerbosity=detailed

Mixed Precision

To keep a specific layer in higher precision while quantizing the rest to INT8, pass its nn.Conv2d instance to ModulesToSkip after transform. Note that torch.fx.GraphModule deep-copies submodules during tracing, so you must take the reference from the fused graph, not from the original model:

from embedl_deploy.quantize import quantize, QuantConfig, ModulesToSkip

res = transform(model, patterns=TENSORRT_PATTERNS)

# Grab the conv instance from the fused graph (not from the original model)
first_conv = res.model.FusedConvBNActMaxPool_0.conv

config = QuantConfig(
    skip=ModulesToSkip(
        stub={first_conv},    # disables input activation quantization
        weight={first_conv},  # disables weight fake-quantization
    )
)
quantized_model = quantize(
    res.model, (example_input,), config=config, forward_loop=calibration_loop
)

Design Principles

  1. Patterns are the only abstraction. Every graph transformation — fusion, conversion, quantization — is a Pattern subclass. Adding a new backend (TIDL, QNN, …) means defining a new set of Pattern subclasses and fused modules with quantization information. The core plan/apply machinery stays the same.

  2. Plans are editable. get_transformation_plan() returns a plan the user can inspect and edit before applying. Toggle match.apply = False to skip specific matches. transform() is a convenience for the common case where you want everything applied.

  3. Graph-based models (torch.export.export and symbolic traced). All graph analysis and surgery uses traced graphs. Models are traced once and manipulated as fx.GraphModule objects with suport for tracing via both torch.fx (symbolic) as well as torch.export.export (Aten). Support for Aten graphs is automatically enabled using Aten recomposition patterns that compose Aten operations into equivalent torch.nn modules automatically before conversions and fusions.

Support

License

Free for non-commercial use within the Embedl Community License (v.1.0).

Please Contact us for commercial licensing.

Copyright (C) 2026 Embedl AB

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

embedl_deploy_tensorrt-0.6.0.tar.gz (56.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

embedl_deploy_tensorrt-0.6.0-py3-none-any.whl (68.1 kB view details)

Uploaded Python 3

File details

Details for the file embedl_deploy_tensorrt-0.6.0.tar.gz.

File metadata

  • Download URL: embedl_deploy_tensorrt-0.6.0.tar.gz
  • Upload date:
  • Size: 56.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.10.20

File hashes

Hashes for embedl_deploy_tensorrt-0.6.0.tar.gz
Algorithm Hash digest
SHA256 7ca4114d4a0821d345c28ffd5e6c0938f1af0a8347bd22256eb60f0803f1c9f6
MD5 ac0d3b785c74d3bb1b0d93200b742b77
BLAKE2b-256 d92c9f1ae116a5334114fd4469b5baea5df3e614d2e7a847dd3c2969722dabc2

See more details on using hashes here.

File details

Details for the file embedl_deploy_tensorrt-0.6.0-py3-none-any.whl.

File metadata

File hashes

Hashes for embedl_deploy_tensorrt-0.6.0-py3-none-any.whl
Algorithm Hash digest
SHA256 9996c6d32ebfa2dd75dad3e9447d25df632753aefa31a9447b1466057f08977e
MD5 4e75aae588766b7037df3bca2162d69c
BLAKE2b-256 c429da310099f415300e17d37fb26615b8e648136a3c6027a10b2754ba7b2ee6

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page