Dynamic inputs for PyTorch TensorRT model?

Hi, I have looked into PyTorch TensorRT document, I have a question in the below line the inputs variable takes min_shape, opt_shape, max_shape does it means that I can leverage this for my use-case where my model takes dynamic input tensors.

import torch_tensorrt
model = MyModel().eval() # torch module needs to be in eval (not training) mode

inputs = [torch_tensorrt.Input(
            min_shape=[1, 1, 16, 16],
            opt_shape=[1, 1, 32, 32],
            max_shape=[1, 1, 64, 64],
            dtype=torch.half,
        )]

I have built the model which takes dynamic inputs like the below example,

1. example1 = torch.rand(1, 3, 224, 224)
2. example2 = torch.rand(1, 3, 224, 550)
3. example3 = torch.rand(1, 3, 640, 480)

so, can I consider example1 as min_shape and example2 as max_shape ?

like

inputs = [torch_tensorrt.Input(
            min_shape=[1, 3, 224, 224],
            max_shape=[1, 3, 640, 480],
            dtype=torch.half,
        )]

could you please clarify this? @ptrblck

Yes, I think you are right as the description fits the docs:

shape_mode (torch_tensorrt.Input._ShapeMode) – Is input statically or dynamically shaped
shape (Tuple or Dict) –
Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form …

1 Like

@ptrblck thanks for confirming :slight_smile: