❓ [Question] How do you use dynamic shape when using fx as ir and the model is not fully lowerable

[ivan94fi]

❓ Question

I have a pytorch model that contains a Pixel Shuffle operation (which is not fully supported) and I would like to convert it to TensorRT, while being able to specify a dynamic shape as input. The "ts" path does not work as there is an issue, the "fx" path has problems too and I am not able to use a splitted model with dynamic shapes.

What you have already tried

• The conversion using TorchScript as "ir" is not working (see Issue 🐛 [Bug] PixelShuffle with dynamic shape error #1568)
• The conversion using torch_tensorrt.fx.compile succeeds when I use a static shape, however there is no way of specifying a dynamic shape
• Using a manual approach (that is by manually tracing with acc_tracer, then constructing the TRTInterpreter and finally the TRTModule) fails as there is a non supported operation (a pixel shuffle layer) (Maybe I should open an Issue for this too?)
• Using the manual approach with a TRTSplitter is maybe the way to go but I don't know how to specify the dynamic shape constraints in this situation.

The "manual" approach that I mentioned is the one specified in examples/fx/fx2trt_example.py and in the docs.

Here is the code as I have it now. Please note that the branch with the splitter is executed and the result is errors when I execute the trt model with different shapes. If do_split is set to False the conversion fails as nn.PixelShuffle is not supported.

import tensorrt as trt
import torch.fx
import torch.nn as nn

import torch_tensorrt.fx.tracer.acc_tracer.acc_tracer as acc_tracer
import torchvision.models as models
from torch_tensorrt.fx import InputTensorSpec, TRTInterpreter, TRTModule
from torch_tensorrt.fx.utils import LowerPrecision
from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter


class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.shuffle = nn.PixelShuffle(2)

    def forward(self, x):
        return self.shuffle(self.conv(x))


torch.set_grad_enabled(False)

# inputs
inputs = [torch.rand(1, 3, 224, 224).cuda()]


factory_kwargs = {"dtype": torch.float32, "device": torch.device("cuda:0")}
model = MyModel().to(**factory_kwargs)

model = model.eval()

out = model(inputs[0])

# sybolic trace
acc_model = acc_tracer.trace(model, inputs)

do_split = True

if do_split:
    # split
    splitter = TRTSplitter(acc_model, inputs)

    splitter.node_support_preview(dump_graph=False)

    split_mod = splitter()

    print(split_mod.graph)

    def get_submod_inputs(mod, submod, inputs):
        acc_inputs = None

        def get_input(self, inputs):
            nonlocal acc_inputs
            acc_inputs = inputs

        handle = submod.register_forward_pre_hook(get_input)
        mod(*inputs)
        handle.remove()
        return acc_inputs

    for name, _ in split_mod.named_children():
        if "_run_on_acc" in name:
            submod = getattr(split_mod, name)
            # Get submodule inputs for fx2trt
            acc_inputs = get_submod_inputs(split_mod, submod, inputs)

            # fx2trt replacement
            interp = TRTInterpreter(
                submod,
                InputTensorSpec.from_tensors(acc_inputs),
                explicit_batch_dimension=True,
            )
            r = interp.run(lower_precision=LowerPrecision.FP32)
            trt_mod = TRTModule(*r)
            setattr(split_mod, name, trt_mod)

    trt_model = split_mod

else:
    # input specs
    input_specs = [
        InputTensorSpec(
            shape=(1, 3, -1, -1),
            dtype=torch.float32,
            device="cuda:0",
            shape_ranges=[((1, 3, 112, 112), (1, 3, 224, 224), (1, 3, 512, 512))],
        ),
    ]
    # input_specs = [
    #     InputTensorSpec(
    #         shape=(1, 3, 224, 224),
    #         dtype=torch.float32,
    #         device="cuda:0",
    #     ),
    # ]

    # TRT interpreter
    interp = TRTInterpreter(
        acc_model,
        input_specs,
        explicit_batch_dimension=True,
        explicit_precision=True,
        logger_level=trt.Logger.INFO,
    )

    interpreter_result = interp.run(
        max_batch_size=4, lower_precision=LowerPrecision.FP32
    )

    # TRT module
    trt_model = TRTModule(
        interpreter_result.engine,
        interpreter_result.input_names,
        interpreter_result.output_names,
    )

trt_out = trt_model(inputs[0])


trt_model(torch.rand(1,3, 112, 112).cuda())
trt_model(torch.rand(1,3, 150, 150).cuda())
trt_model(torch.rand(1,3, 400, 400).cuda())
trt_model(torch.rand(1,3, 512, 512).cuda())

print((trt_out - out).max())

Environment

The official NVIDIA Pytorch Docker image version 22.12 is used.

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0): 1.3.0a0
• PyTorch Version (e.g. 1.0): 1.14.0a0+410ce96
• CPU Architecture: AMD64
• OS (e.g., Linux): Ubuntu
• How you installed PyTorch (conda, pip, libtorch, source): preinstalled in the Docker image
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version: 3.8.10
• CUDA version: 12.0
• GPU models and configuration: NVIDIA GeForce RTX 2080 Ti, driver version 525.60.11
• Any other relevant information:

[github-actions]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days