Graph Conversions#

Conversions are structural graph rewrites that normalize the computation graph before fusion patterns are applied. They change the graph topology — replacing one set of operators with a functionally equivalent set that downstream patterns can match.

Conversion patterns have phase = Phase.CONVERSION and are applied iteratively by transform() until no new matches are found.

Built-in TensorRT conversions#

FlattenLinearToConv1x1Pattern#

Matches: Flatten (4D→2D) → [Dropout/Activation]* → Linear

Replaces with: [Dropout/Activation]* → Conv2d(1×1) → Flatten

Many classification networks end with:

AdaptiveAvgPool2d → Flatten → Linear(1000)

The Linear layer cannot be fused with preceding operations in TensorRT. After this conversion it becomes:

AdaptiveAvgPool2d → Conv2d(1×1, out=1000) → Flatten

The Conv2d(1×1) can then be matched by downstream fusion patterns (ConvBNPattern, etc.) and benefits from INT8 quantization.

Weight conversion: The Linear weight matrix of shape (out_features, in_features) is reshaped to (out_features, in_features, 1, 1) for the equivalent Conv2d.

Element-wise ops: Activations or dropout layers between Flatten and Linear are absorbed by a Wildcard and moved before the Conv2d in the replacement graph.

Affected architectures: ResNet, ConvNeXt, EfficientNet, MobileNet, and most classification backbones with a Flatten → Linear classifier head.

RemoveIdentityAdaptiveAvgPoolPattern#

Matches: AdaptiveAvgPool2d where output_size == input spatial dims

Replaces with: nothing (erases the node)

In ConvNeXt-style architectures, AdaptiveAvgPool2d(output_size=(7, 7)) is applied to a 7×7 feature map — a mathematical identity. Removing it simplifies the graph and prevents it from blocking fusion of surrounding operators.

Note

Requires shape metadata (via torch.fx.passes.shape_prop.ShapeProp) to determine that input and output spatial dimensions match. Nodes without shape metadata are skipped with a warning.

RemoveIdentityPattern#

Matches: nn.Identity

Replaces with: nothing (erases the node)

Some architectures include explicit nn.Identity modules as placeholders (e.g. for optional dropout). Removing them simplifies the graph so surrounding operators can be matched by fusion patterns.

RemoveAssertPattern#

Matches: getattr/size/dim → eq → torch._assert

Replaces with: nothing (erases the nodes)

Export-traced graphs contain assertion chains that verify tensor shapes at runtime. These are noise for the fusion pass — removing them prevents fragmented matches.

DecomposeMultiheadAttentionPattern#

Matches: nn.MultiheadAttention (self-attention, batch_first=True)

Replaces with: three explicit modules:

  1. MHAInProjection — the combined Q/K/V linear projection

  2. ScaledDotProductAttention — the attention computation

  3. nn.Linear — the output projection

PyTorch’s nn.MultiheadAttention is a monolithic module. TensorRT cannot fuse or quantize its internal operations. By decomposing it into visible sub-modules in the FX graph, each component can be independently fused and quantized.

Restrictions: Only self-attention (_qkv_same_embed_dim=True, batch_first=True) without masks or add_zero_attn is supported. Unsupported configurations are skipped with a warning.

Affected architectures: Vision Transformer (ViT), DeiT, and any model using nn.MultiheadAttention.

ComposeScaledDotProductAttentionPattern#

Matches: Q → matmul → scale → softmax → [dropout] → matmul(V) chain

Replaces with: ScaledDotProductAttention module

When nn.MultiheadAttention is traced via torch.export, the attention computation is lowered to individual matmul/softmax/scale ops. This conversion recomposes them into a single ScaledDotProductAttention module that the fusion pass can wrap.

WrapFunctionalSDPAPattern#

Matches: F.scaled_dot_product_attention call

Replaces with: ScaledDotProductAttention module

Models that call the functional SDPA API directly (rather than via nn.MultiheadAttention) produce a call_function node. This conversion wraps it in the same ScaledDotProductAttention module used by ComposeScaledDotProductAttentionPattern.

ComposeParallelLinearsPattern#

Matches: three parallel transpose → view → Linear branches from the same source into a ScaledDotProductAttention

Replaces with: MHAInProjection (packs Q/K/V into one Linear(embed_dim, 3 * embed_dim))

After attention decomposition, the Q/K/V projections are separate linear layers. This conversion fuses them into a single projection for efficient quantization and export.

DecomposeSwinAttentionPattern#

Matches: torchvision.models.swin_transformer.shifted_window_attention function call

Replaces with: a sequence of explicit modules:

  1. SwinWindowPartition — partition input into windows

  2. MHAInProjection — Q/K/V projection

  3. SwinAttention — windowed attention computation

  4. nn.Linear — output projection

  5. SwinWindowReverse — reassemble windows into the feature map

Like DecomposeMultiheadAttentionPattern, this makes the internal operations visible for independent fusion and quantization. Requires torchvision to be installed; skipped with a warning otherwise.

Affected architectures: Swin Transformer V1/V2.

When conversions matter#

ResNet50#

The FlattenLinearToConv1x1Pattern converts the final classifier:

avgpool (AdaptiveAvgPool2d) → flatten → fc (Linear)

becomes:

avgpool (AdaptiveAvgPool2d) → fc (Conv2d 1×1) → flatten

The Conv2d is then picked up by ConvBNPattern during the fusion stage.

ConvNeXt#

ConvNeXt uses AdaptiveAvgPool2d at several points where the spatial dimensions do not change. The RemoveIdentityAdaptiveAvgPoolPattern cleans these up, and FlattenLinearToConv1x1Pattern converts the head.

For ConvNeXt models, shape propagation is required before running conversions. transform() handles this automatically:

example_input = torch.randn(1, 3, 224, 224)
result = transform(model, (example_input,))

Vision Transformer (ViT)#

DecomposeMultiheadAttentionPattern expands each nn.MultiheadAttention into three sub-modules. After decomposition, the fusion pass applies MHAInProjectionPattern and ScaledDotProductAttentionPattern to wrap them in quantization-aware fused modules. The output projection nn.Linear is matched by LinearPattern.

Running conversions only#

from embedl_deploy import transform
from embedl_deploy.tensorrt import TENSORRT_CONVERSION_PATTERNS

# Apply only structural conversions
result = transform(model, (example_input,), patterns=TENSORRT_CONVERSION_PATTERNS)
converted_model = result.model

This is useful for debugging: you can inspect the graph after conversions to verify the structural rewrites before running fusions.