Quickstart#
Embedl Deploy keeps deployment preparation inside PyTorch.
You provide a model and a hardware-target pattern set, and
transform() rewrites the graph into a deployable form.
Core concepts#
Modules: PyTorch
nn.Moduleobjects (e.g.,ConvBatchNormRelu).Patterns: Hardware-aware graph patterns such as fusions, conversions and quantization placements. A pattern finds a subgraph and replaces it with a hardware-aware equivalent. Patterns are carefully determined by extensive exploration of the hardware toolchain, restrictions and verified
Transform: applies patterns to your model in one pass.
The public release contains TensorRT patterns, within the following groups:
Fusion patterns: operator fusions, e.g.,
Conv+BatchNorm+ReLUConversion patterns: structural rewrites replacing one layer with another
Quantized patterns: quantization-focused rewrites
The only requirement to run a transform is to pass a pattern set for your hardware target.
Usage#
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
)
Visualization and logging on Embedl Hub#
The Embedl visualizer is a graphical tool
that can be used to render PyTorch graphs, ONNX models, and hardware-compiled
artifacts (e.g., TensorRT compiled engines, TIDL models) side-by-side for
comparison and debugging. It is available online for public use on the Embedl
Hub and locally for enterprise solutions. In order to
log metrics and models in the Hub, follow the simple code snippet below after
creating an Embedl hub account, installing the hub client with
pip install embedl-hub and logging to the hub with an API key created on the
hub.
embedl-hub auth --api-key <YOUR_API_KEY>
Then you can run experiments, log artifacts and use the visualizer with the following simple snippet.
import torch
from torchvision.models import resnet18 as Model
from embedl_deploy import transform
from embedl_deploy.quantize import quantize
from embedl_deploy.tensorrt import TENSORRT_PATTERNS
from embedl_hub.tracking import Client
client = Client()
client.set_project("Deploy Demo")
# 1. Load a standard PyTorch model.
model = Model().eval()
example_input = torch.randn(1, 3, 224, 224)
with client.start_run("compile", name="ResNet18 TensorRT PTQ"):
client.log_param("model", "resnet18")
# 2. Transform: fuse and optimize for TensorRT in one call.
# For more compatibility you can trace your model with
# torch.export.export:
# 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]))
client.log_metric("transform_matches", len(res.matches))
# 3. Quantize with post-training quantization.
def calibration_loop(model: torch.fx.GraphModule) -> None:
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. Save the QDQ graph as a pt2 file for later reuse.
pt2_path = "model_qdq.pt2"
exported = torch.export.export(quantized_model, (example_input,))
torch.export.save(exported, pt2_path)
client.log_artifact(pt2_path, name="model_qdq_pt2")
# 5. Export as usual. Dynamo-exported models may have compilation
# issues in some environments.
onnx_path = "model.onnx"
torch.onnx.export(
quantized_model,
(example_input,),
onnx_path,
dynamo=False,
)
client.log_artifact(onnx_path, name="model_onnx")
The images below show the layer mapping before and after TensorRT compilation.