torch2trt converters to bootstrap FX converter library

[apbose]

torch2trt converters to bootstrap FX converter library

TL;DR

Bootstrap the FX frontend converter library with the Torch Script writing aten -> TRT converters in torch2trt converters https://github.com/NVIDIA-AI-IOT/torch2trt

Goal(s)

The converter library in torch2trt has more operations support as compared to fx2trt. The goal is to leverage this and incorporate it in the fx2trt design flow. The primary goal of this document is highlight the next steps - API design, code flow, and the use cases to enhance the present converter library in fx2trt

Usecases

Currently fx2trt supports the following converters- adaptive_avgpool, activation, add, batchnorm, convolution, div, linear, maxpool, mul,pow, relu, sub and reshape operations. torch2trt has a bunch of converters which are registered using the tensorrt_converter decorater in the same way as the converters are registered in fx2trt. Goal is to support the missing operations and create a generic interface for fx converters which will then be used for the aten2trt interface

Proposed APIs / UX

In torch2trt the converters are registered in the CONVERTERS dictionary. The conversion context in torch2trt and fx2trt behave in the same way. It attaches the converters to the torch method calls. Currently the acc_ops have the same syntax as aten ops due to which they are reused.

Design

The tensorrt_converter object attaches the converter to the torch method calls in the trace. Aten and acc converters have different protorypes. A common conversion context should wrap up inputs and arguments in such a way, that the interface should be common to the acc and aten tracer.

Implementation Phases

1. Run the unit tests in TensorRT/py/torch_tensorrt/fx/test/aten_op/*.py to study through the aten trace (currently in progress)
2. For including the ops, the flow is -

• If 1-1 mapping exists. Just remap the implementation to acc_ops_converters
• If it does not exist, we can create corresponding implementation in aten_ops_converters.py

3. Run progressively the model list and start supporting all the ops in the models. The model list:

• squeezenet1_0
• mobilenet_v2
• inception_v3
• efficientnet_b0
• regnet_y_8gf

4. In the process, integrate the converters such that they are agnostic to the aten and acc tracers. So we would want to make the flow generic to fx2trt, instead of seperate as illustrated in the second point
   For this the following changes are to be made-
   -The input signature for the converters in aten and acc is same, but the way arguments are packed is different.
   eg:

Aten trace

.graph():
    %arg0 : [#users=1] = placeholder[target=arg0]
    %_adaptive_avg_pool2d_default : [#users=1] = call_function[target=torch.ops.aten._adaptive_avg_pool2d.default](args = (%arg0, [64, 64]), kwargs = {})
    return [_adaptive_avg_pool2d_default]

Acc trace

graph():
    %x : [#users=1] = placeholder[target=x]
    %adaptive_avg_pool2d_1 : [#users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.adaptive_avg_pool2d](args = (), kwargs = {input: %x, output_size: (64, 64)})
    return adaptive_avg_pool2d_1

In the above example fx2trt can maintain a common Conversion context in which the arguments can be passed for Aten and Acc Tracer
For the missing converters, the converters from torch2trt can be used directly

Prototype - M

• Demonstrate process porting torch2trt converters to FX

MVP (<TARGET RELEASE VERSION>) - S

• Implement 5 high impact converters

[frank-wei]

Thanks for the RFC @apbose . A few comments here:

1. Currently, fx2trt has supported quite a lot of the ops that common models used and tested widely in our internal prod scenarios. We can look for the missing ones from torch2trt to better support the op coverage. In another words, we can reuse the existing fx converters which save us a lot of efforts.
2. The ConversionContext should behave similar as the interpreter to traverse the graph nodes and make func call for the conversion. fx2trt has it called TRTInterpreter and defined here https://github.com/pytorch/TensorRT/blob/master/py/torch_tensorrt/fx/fx2trt.py
3. I wrote some preliminary op conversion in https://github.com/pytorch/TensorRT/tree/master/py/torch_tensorrt/fx/test/converters/aten_op It is a good start for us to begin support the aten2trt (aten is the PT2.0 tracer IR)
   Aten2trt is our next gen of converter which connects PT2.0 and TensorRT. We hope it can reuse the current infra to set up the workflow quickly and be able to verify our concept. I will also start to work with you to initiate this integration.
   cc @ncomly-nvidia @narendasan @yinghai

[ncomly-nvidia]

@frank-wei do you see benefit for appending these ops or will they be overshadowed by aten in once 2.0 is released?

[frank-wei]

From converter perspective, most acc_ops has the same syntax as aten ops and that is why we can reuse most of them. For example, these ops link link ...etc.

[yinghai]

Yeah I think it'd be good if we consolidate the FX based conversion effort (including pt 2.0) to the pytorch/TensorRT codebase where we already have support for a lot of of the pytorch ops and they are being extensively used/tested.

[ncomly-nvidia]

@apbose @narendasan can you confirm here if this will be still useful after proxy_tensor & aten ops are the path in FX?

[SrivastavaKshitij]

@ncomly-nvidia :we have another 20 or so converters that can be open-sourced to support this effort. Some of those converters support dynamic shapes. Though that means editing torch2trt.py file as well but that's not a big lift.