Prototype reorg of the FX converters to clean dependency chain #1683
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…so that the three tracing paths call down into a common converter base instead of across each other
narendasan
added
the
WIP
github-actions
bot
added
component: api [Python]
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- py/torch_tensorrt/fx/converters/activation.py 2023-02-20 22:37:19.697378 +0000
+++ py/torch_tensorrt/fx/converters/activation.py 2023-02-20 22:37:34.727752 +0000
@@ -19,19 +19,20 @@
TRTPluginFieldCollection,
TRTTensor,
)
from ..utils import torch_dtype_from_trt
+
def add_activation_layer(
network: TRTNetwork,
input_val: TRTTensor,
operation_type: trt.ActivationType,
target: Target,
name: str,
alpha: Optional[Any] = None,
beta: Optional[Any] = None,
- dyn_range_fn: Optional[Callable[Tuple[float, float]]] = None
+ dyn_range_fn: Optional[Callable[Tuple[float, float]]] = None,
) -> TRTTensor:
"""
Add a TensorRT Activation layer to `network`.
Args:
@@ -68,20 +69,22 @@
dyn_range = dyn_range_fn(input_val.dynamic_range)
mark_as_int8_layer(layer, dyn_range)
return layer.get_output(0)
+
def add_elu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
) -> TRTTensor:
input_val = kwargs["input"]
alpha = kwargs["alpha"]
operation_type = trt.ActivationType.ELU
return add_activation_layer(network, input_val, operation_type, target, name, alpha)
+
def add_gelu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
@@ -113,10 +116,11 @@
layer = network.add_plugin_v2([input_val], plugin)
set_layer_name(layer, target, name)
return layer.get_output(0)
+
def add_hard_sigmoid(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -130,10 +134,11 @@
name,
alpha=1 / 6,
beta=0.5,
)
+
def add_hardtanh(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -168,10 +173,11 @@
operation_type = trt.ActivationType.LEAKY_RELU
return add_activation_layer(
network, input_val, operation_type, target, name, negative_slope
)
+
def add_relu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -180,21 +186,30 @@
operation_type = trt.ActivationType.RELU
def activation_dyn_range_fn(dyn_range):
return max(0, dyn_range[0]), max(0, dyn_range[1])
- return add_activation_layer(network, input_val, operation_type, target, name, dyn_range_fn=activation_dyn_range_fn)
+ return add_activation_layer(
+ network,
+ input_val,
+ operation_type,
+ target,
+ name,
+ dyn_range_fn=activation_dyn_range_fn,
+ )
+
def add_selu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
) -> TRTTensor:
input_val = kwargs["input"]
operation_type = trt.ActivationType.SELU
return add_activation_layer(network, input_val, operation_type, target, name)
+
def add_sigmoid(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
@@ -207,11 +222,16 @@
return 1 / (1 + np.exp(-x))
return sigmoid_fn(dyn_range[0]), sigmoid_fn(dyn_range[1])
return add_activation_layer(
- network, input_val, trt.ActivationType.SIGMOID, target, name, dyn_range_fn=activation_dyn_range_fn
+ network,
+ input_val,
+ trt.ActivationType.SIGMOID,
+ target,
+ name,
+ dyn_range_fn=activation_dyn_range_fn,
)
def add_softsign(
network: TRTNetwork,
@@ -221,10 +241,11 @@
) -> TRTTensor:
input_val = kwargs["input"]
operation_type = trt.ActivationType.SOFTSIGN
return add_activation_layer(network, input_val, operation_type, target, name)
+
def add_tanh(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
--- py/torch_tensorrt/fx/converters/nn_ops_converters.py 2023-02-20 22:37:19.697378 +0000
+++ py/torch_tensorrt/fx/converters/nn_ops_converters.py 2023-02-20 22:37:34.990162 +0000
@@ -7,17 +7,19 @@
from ..converter_registry import tensorrt_converter
from .converter_utils import mark_as_int8_layer
import activation
+
@tensorrt_converter(torch.nn.functional.relu)
@tensorrt_converter(torch.nn.modules.activation.ReLU)
def relu(network, submod, args, kwargs, layer_name):
# args/kwargs should have already been normalized to kwargs
assert len(args) == 0
- return activation.add_relu(network,"tensorrt", kwargs, layer_name)
+ return activation.add_relu(network, "tensorrt", kwargs, layer_name)
+
@tensorrt_converter(torch.nn.modules.activation.Sigmoid)
def sigmoid(network, submod, args, kwargs, layer_name):
# args/kwargs should have already been normalized to kwargs
assert len(args) == 0
- return activation.add_sigmoid(network,"tensorrt", kwargs, layer_name)
\ No newline at end of file
+ return activation.add_sigmoid(network, "tensorrt", kwargs, layer_name)
--- py/torch_tensorrt/fx/converters/acc_ops_converters.py 2023-02-20 22:37:19.697378 +0000
+++ py/torch_tensorrt/fx/converters/acc_ops_converters.py 2023-02-20 22:37:36.744356 +0000
@@ -1027,10 +1027,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_elu(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.selu)
def acc_ops_selu(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -1048,10 +1049,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_softsign(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.sin)
def acc_ops_sin(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -1122,10 +1124,11 @@
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_tanh(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.asin)
def acc_ops_asin(
network: TRTNetwork,
target: Target,
@@ -3323,10 +3326,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_gelu(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.chunk)
def acc_ops_chunk(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -3477,10 +3481,11 @@
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_hardtanh(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.interpolate)
def acc_ops_interpolate(
network: TRTNetwork,
target: Target,
There was a problem hiding this comment.
There are some changes that do not conform to Python style guidelines:
--- py/torch_tensorrt/fx/converters/activation.py 2023-02-20 22:37:19.800788 +0000
+++ py/torch_tensorrt/fx/converters/activation.py 2023-02-20 22:37:37.322661 +0000
@@ -19,19 +19,20 @@
TRTPluginFieldCollection,
TRTTensor,
)
from ..utils import torch_dtype_from_trt
+
def add_activation_layer(
network: TRTNetwork,
input_val: TRTTensor,
operation_type: trt.ActivationType,
target: Target,
name: str,
alpha: Optional[Any] = None,
beta: Optional[Any] = None,
- dyn_range_fn: Optional[Callable[Tuple[float, float]]] = None
+ dyn_range_fn: Optional[Callable[Tuple[float, float]]] = None,
) -> TRTTensor:
"""
Add a TensorRT Activation layer to `network`.
Args:
@@ -68,20 +69,22 @@
dyn_range = dyn_range_fn(input_val.dynamic_range)
mark_as_int8_layer(layer, dyn_range)
return layer.get_output(0)
+
def add_elu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
) -> TRTTensor:
input_val = kwargs["input"]
alpha = kwargs["alpha"]
operation_type = trt.ActivationType.ELU
return add_activation_layer(network, input_val, operation_type, target, name, alpha)
+
def add_gelu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
@@ -113,10 +116,11 @@
layer = network.add_plugin_v2([input_val], plugin)
set_layer_name(layer, target, name)
return layer.get_output(0)
+
def add_hard_sigmoid(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -130,10 +134,11 @@
name,
alpha=1 / 6,
beta=0.5,
)
+
def add_hardtanh(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -168,10 +173,11 @@
operation_type = trt.ActivationType.LEAKY_RELU
return add_activation_layer(
network, input_val, operation_type, target, name, negative_slope
)
+
def add_relu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
@@ -180,21 +186,30 @@
operation_type = trt.ActivationType.RELU
def activation_dyn_range_fn(dyn_range):
return max(0, dyn_range[0]), max(0, dyn_range[1])
- return add_activation_layer(network, input_val, operation_type, target, name, dyn_range_fn=activation_dyn_range_fn)
+ return add_activation_layer(
+ network,
+ input_val,
+ operation_type,
+ target,
+ name,
+ dyn_range_fn=activation_dyn_range_fn,
+ )
+
def add_selu(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
) -> TRTTensor:
input_val = kwargs["input"]
operation_type = trt.ActivationType.SELU
return add_activation_layer(network, input_val, operation_type, target, name)
+
def add_sigmoid(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
@@ -207,11 +222,16 @@
return 1 / (1 + np.exp(-x))
return sigmoid_fn(dyn_range[0]), sigmoid_fn(dyn_range[1])
return add_activation_layer(
- network, input_val, trt.ActivationType.SIGMOID, target, name, dyn_range_fn=activation_dyn_range_fn
+ network,
+ input_val,
+ trt.ActivationType.SIGMOID,
+ target,
+ name,
+ dyn_range_fn=activation_dyn_range_fn,
)
def add_softsign(
network: TRTNetwork,
@@ -221,10 +241,11 @@
) -> TRTTensor:
input_val = kwargs["input"]
operation_type = trt.ActivationType.SOFTSIGN
return add_activation_layer(network, input_val, operation_type, target, name)
+
def add_tanh(
network: TRTNetwork,
target: Target,
kwargs: Dict[str, Argument],
name: str,
--- py/torch_tensorrt/fx/converters/nn_ops_converters.py 2023-02-20 22:37:19.800788 +0000
+++ py/torch_tensorrt/fx/converters/nn_ops_converters.py 2023-02-20 22:37:37.627636 +0000
@@ -7,17 +7,19 @@
from ..converter_registry import tensorrt_converter
from .converter_utils import mark_as_int8_layer
import activation
+
@tensorrt_converter(torch.nn.functional.relu)
@tensorrt_converter(torch.nn.modules.activation.ReLU)
def relu(network, submod, args, kwargs, layer_name):
# args/kwargs should have already been normalized to kwargs
assert len(args) == 0
- return activation.add_relu(network,"tensorrt", kwargs, layer_name)
+ return activation.add_relu(network, "tensorrt", kwargs, layer_name)
+
@tensorrt_converter(torch.nn.modules.activation.Sigmoid)
def sigmoid(network, submod, args, kwargs, layer_name):
# args/kwargs should have already been normalized to kwargs
assert len(args) == 0
- return activation.add_sigmoid(network,"tensorrt", kwargs, layer_name)
\ No newline at end of file
+ return activation.add_sigmoid(network, "tensorrt", kwargs, layer_name)
--- py/torch_tensorrt/fx/converters/acc_ops_converters.py 2023-02-20 22:37:19.800788 +0000
+++ py/torch_tensorrt/fx/converters/acc_ops_converters.py 2023-02-20 22:37:39.910005 +0000
@@ -1027,10 +1027,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_elu(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.selu)
def acc_ops_selu(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -1048,10 +1049,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_softsign(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.sin)
def acc_ops_sin(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -1122,10 +1124,11 @@
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_tanh(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.asin)
def acc_ops_asin(
network: TRTNetwork,
target: Target,
@@ -3323,10 +3326,11 @@
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_gelu(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.chunk)
def acc_ops_chunk(
network: TRTNetwork,
target: Target,
args: Tuple[Argument, ...],
@@ -3477,10 +3481,11 @@
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return activation.add_hardtanh(network, target, kwargs, name)
+
@tensorrt_converter(acc_ops.interpolate)
def acc_ops_interpolate(
network: TRTNetwork,
target: Target,
|
@apbose This is sort of what I have been talking about. We abstract out the various potential IRs in files like |
apbose
approved these changes
Feb 24, 2023
apbose
reviewed
Feb 28, 2023
|
Raised PR #1745 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR proposes a reorg of the converter library for FX so that the dependency chain is much cleaner considering there's 3 versions of ops being supported in FX (nn, acc, aten).
Instead of having implementations call each other, the idea is to have an IR agnostic set of converters and then the
[aten/acc/nn]_ops_converterfiles will just pack arguments properly.Fixes # (issue)
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: