aten::size with dynamic shapes handling

[peri044]

Goal:

We currently do not support models with dynamic shapes which have aten::size in them.

Two relevant issues:

1. 🐛 [Bug] Compilation failure for SSD300 model with dynamic batch #1555 (SSD model)
   In this model, aten::size is the input to aten::reshape layer. For a dynamic shaped input, aten::size just outputs -1 instead of a shape tensor and hence shape information is not propagated down the network resulting in errors.
2. Resnet model with dynamic shapes
   Torchvision resnet passes with dynamic shapes. However, there is an alternate implementation by Nvidia (used by the model navigator) that fails.
   Torchvision implementation of final avg pool and FC layers: https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py#L279
   Nvidia's implementation: https://github.com/NVIDIA/DeepLearningExamples/blob/c2bb3fea797403612a5ea8e359eb31e7e750374f/PyTorch/Classification/ConvNets/image_classification/models/resnet.py#L312
   The latter implementation explicitly uses size() call which doesn't work with dynamic shapes.

Both of these issues use dynamic batch and not dynamic shapes.
Both these issues fail using Torchscript backend. However, with some tweaks, FX backend compilation is successful.
FX reshape converters have the support for accepting ITensors as a second input.

The main change required for successful compilation of the above models is to add dynamic shape support in FX.

Dynamic shape support in FX:
Currently FX has an option dynamic_batch=True. But the support is not fully implemented. https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/fx/input_tensor_spec.py#L21-L29 hardcodes the batch size ranges for dynamic batch inputs.

1. For dynamic batch case, if we modify the https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/fx/input_tensor_spec.py#L21-L29 to provide min, opt, max batches, 🐛 [Bug] Compilation failure for SSD300 model with dynamic batch #1555 runs fine.

2. For dynamic shapes, verify the toy positional embedding bug which has dynamic shape inputs. Torchscript backend fails but verify it in FX.

Proposal:

Our current Input class uses cpp code via ._C binding.

from torch_tensorrt import _C
class Input(object):
  def __init__: ...
  def _to_internal(self) -> _C.Input:
  def _supported_input_size_type(input_size: Any) -> bool:
  def _parse_dtype(dtype: Any) -> _enums.dtype:
  def from_tensor(cls, t: torch.Tensor) -> "Input":
  def from_tensors(cls, ts: torch.Tensor) -> List["Input"]:
  def example_tensor(self, optimization_profile_field: str = None) -> torch.Tensor 

This is called as torch_tensorrt.Input() and passed to torch_tensorrt.compile

On torchscript side, we convert this Input into a _C.Input() by calling _to_internal from the compile_spec.py

1. Make torch_tensorrt.Input compatible with FX. Use the same constructors and methods as the current Torchscript front end. This allows us to maintain consistency of how dynamic shaped inputs can be provided between the two backends.

2. On the torchscript side, in compile_spec.py, we could define a new torch_tensorrt.ts.Input() which consumes this torch_tensorrt.Input and has conversions to _C.Input.

3. On the FX side, we already have a InputTensorSpec object which defines a input spec for a model. Replace this to use torch_tensorrt.Input class directly which should also enable dynamic batch support as well.

4. torch_tensorrt.Input class should be pure pythonic. If users install fx-only path, we should ensure _C.Input() dependency is not involved.

5. Although our current torch_tensorrt.Input class is flexible to accept both dynamic shapes and dynamic batch, we should issue a warning to users using FX path that dynamic shapes are not supported.

Milestones:

MVP: (S-M)

• Make torch_tensorrt.Input compatible with FX.
• Add support for dynamic batch in FX

MVP: (S-M)

• Add aten::size + dynamic shape support in Torchscript

Phase 2: (S-M)

• Add support for dynamic shapes in FX
• Investigate positional embedding issue after dynamic shape support in FX.

[frank-wei]

yes, FX path only supports dynamic batch but not dynamic shape. That is the dilemma of it since trace could not provide this information.

[ncomly-nvidia]

Can proxy tensor also not provide this info?

[frank-wei]

yes, PT2.0 tracer provides the shape information which should be perfect solution from trace perspective. But not sure how much overhead it brings in shape operation. A good example is shown below. aten.sym_size extracts dim=1 and used for the following instruction.

%sym_size : [#users=1] = call_function[target=torch.ops.aten.sym_size](args = (%slice_174, 1), kwargs = {})
%add : [#users=1] = call_function[target=operator.add](args = (%sym_size, 4), kwargs = {})

But this instruction aten.sym_size needs to convert to a shape layer in TRT and I am not sure if it becomes a big overhead once similar instruction becomes more in one model.