Converter reorg batch norm

batch norm error fix and linting issue error fix

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -29,6 +29,7 @@
 from torch_tensorrt.fx.converters.impl import activation
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
+from torch_tensorrt.fx.converters.impl.normalization import batch_norm
 from torch_tensorrt.fx.converters.impl.unary import sign
 from torch_tensorrt.fx.converters.impl.elementwise.base import (
     convert_binary_elementwise,
@@ -630,58 +631,20 @@ def acc_ops_batch_norm(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"BatchNorm2d received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    if has_dynamic_shape(input_val.shape):
-        assert input_val.shape[1] != -1, "Channel dim can't be dynamic for batch norm."
-
-    scale = cast(
-        torch.Tensor, to_numpy(cast(torch.Tensor, kwargs["weight"]))
-    ) / np.sqrt(
-        cast(torch.Tensor, to_numpy(cast(torch.Tensor, kwargs["running_var"])))
-        + cast(float, kwargs["eps"])
-    )
-
-    bias = (
-        to_numpy(cast(torch.Tensor, kwargs["bias"]))
-        - to_numpy(cast(torch.Tensor, kwargs["running_mean"])) * scale
+    return batch_norm(
+        network,
+        target,
+        SourceIR.ACC,
+        name,
+        kwargs["input"],
+        kwargs["weight"],
+        kwargs["bias"],
+        kwargs["running_mean"],
+        kwargs["running_var"],
+        kwargs["training"],
+        kwargs["momentum"],
+        kwargs["eps"],
     )
-    power = np.ones_like(scale)
-
-    # For BatchNorm1d, reshape 1d to 2d
-    output_shape = input_val.shape
-    if not network.has_implicit_batch_dimension and len(input_val.shape) < 4:
-        assert (
-            len(get_dynamic_dims(input_val.shape)) <= 1
-        ), "BatchNorm1D with more than one dynamic dims is not currently supported."
-        reshape_layer = network.add_shuffle(input_val)
-        if len(input_val.shape) == 2:
-            reshape_layer.reshape_dims = (input_val.shape[0], input_val.shape[1], 1, 1)
-        else:  # len(input_val.shape) == 3
-            reshape_layer.reshape_dims = (
-                input_val.shape[0],
-                input_val.shape[1],
-                input_val.shape[2],
-                1,
-            )
-        set_layer_name(reshape_layer, target, f"{name}_reshape_2d")
-        input_val = reshape_layer.get_output(0)
-    layer = network.add_scale(input_val, trt.ScaleMode.CHANNEL, bias, scale, power)
-    set_layer_name(layer, target, name)
-
-    # For BatchNorm1d, reshape output back to 1d
-    if not network.has_implicit_batch_dimension and len(output_shape) < 4:
-        reshape_output_layer = network.add_shuffle(layer.get_output(0))
-        reshape_output_layer.reshape_dims = tuple(output_shape)
-        set_layer_name(reshape_output_layer, target, f"{name}_reshape_1d")
-        layer = reshape_output_layer
-    return layer.get_output(0)
 
 
 @tensorrt_converter(acc_ops.layer_norm)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -25,6 +25,7 @@
 from torch_tensorrt.fx.converters.impl.elementwise import rsqrt
 from torch_tensorrt.fx.converters.impl.elementwise import fmod
 from torch_tensorrt.fx.converters.impl.elementwise import rsub
+from torch_tensorrt.fx.converters.impl.normalization import batch_norm
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -93,18 +94,19 @@ def aten_ops_batch_norm(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    kwargs_new = {
-        "input": args[0],
-        "weight": args[1],
-        "bias": args[2],
-        "running_mean": args[3],
-        "running_var": args[4],
-        "training": args[5],
-        "momentum": args[6],
-        "eps": args[7],
-    }
-    return acc_ops_converters.acc_ops_batch_norm(
-        network, target, None, kwargs_new, name
+    return batch_norm(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+        args[2],
+        args[3],
+        args[4],
+        args[5],
+        args[6],
+        args[7],
     )
 
 

py/torch_tensorrt/fx/converters/impl/normalization/__init__.py

@@ -0,0 +1 @@
+from .ops import *

py/torch_tensorrt/fx/converters/impl/normalization/ops.py

@@ -0,0 +1,95 @@
+import operator
+import warnings
+from typing import cast, Union, Callable, Any, Optional, Sequence
+import logging
+
+import numpy as np
+
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+from torch_tensorrt.fx.utils import get_dynamic_dims
+
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    set_layer_name,
+    has_dynamic_shape,
+    to_numpy,
+)
+
+from torch_tensorrt.fx.converters.impl.unary.base import (
+    convert_unary,
+)
+
+from torch_tensorrt.fx.converters.impl.elementwise.base import (
+    convert_binary_elementwise,
+)
+
+_LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def batch_norm(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    running_mean: torch.Tensor,
+    running_var: torch.Tensor,
+    training: torch.Tensor,
+    momentum: torch.Tensor,
+    eps: list,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"BatchNorm2d received input {input} that is not part "
+            "of the TensorRT region!"
+        )
+
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for batch norm."
+
+    scale = cast(torch.Tensor, to_numpy(cast(torch.Tensor, weight))) / np.sqrt(
+        cast(torch.Tensor, to_numpy(cast(torch.Tensor, running_var))) + cast(float, eps)
+    )
+
+    bias = (
+        to_numpy(cast(torch.Tensor, bias))
+        - to_numpy(cast(torch.Tensor, running_mean)) * scale
+    )
+    power = np.ones_like(scale)
+
+    # For BatchNorm1d, reshape 1d to 2d
+    output_shape = input.shape
+    if not network.has_implicit_batch_dimension and len(input.shape) < 4:
+        assert (
+            len(get_dynamic_dims(input.shape)) <= 1
+        ), "BatchNorm1D with more than one dynamic dims is not currently supported."
+        reshape_layer = network.add_shuffle(input)
+        if len(input.shape) == 2:
+            reshape_layer.reshape_dims = (input.shape[0], input.shape[1], 1, 1)
+        else:  # len(input_val.shape) == 3
+            reshape_layer.reshape_dims = (
+                input.shape[0],
+                input.shape[1],
+                input.shape[2],
+                1,
+            )
+        set_layer_name(reshape_layer, target, f"{name}_reshape_2d")
+        input = reshape_layer.get_output(0)
+    layer = network.add_scale(input, trt.ScaleMode.CHANNEL, bias, scale, power)
+    set_layer_name(layer, target, name)
+
+    # For BatchNorm1d, reshape output back to 1d
+    if not network.has_implicit_batch_dimension and len(output_shape) < 4:
+        reshape_output_layer = network.add_shuffle(layer.get_output(0))
+        reshape_output_layer.reshape_dims = tuple(output_shape)
+        set_layer_name(reshape_output_layer, target, f"{name}_reshape_1d")
+        layer = reshape_output_layer
+    return layer.get_output(0)