fix bugs

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -18,6 +18,7 @@ def args_bounds_check(
     return args[i] if len(args) > i else replacement
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.native_batch_norm.default)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.batch_norm)  # type: ignore[misc]
 def aten_ops_batch_norm(
     network: TRTNetwork,
@@ -32,17 +33,18 @@ def aten_ops_batch_norm(
         SourceIR.ATEN,
         name,
         input=args[0],
-        weight=args_bounds_check(args, 1, replacement=1),
-        bias=args_bounds_check(args, 2, replacement=0),
-        running_mean=args_bounds_check(args, 3),
-        running_var=args_bounds_check(args, 4),
-        training=args_bounds_check(args, 5),
-        momentum=args_bounds_check(args, 6, replacement=0.1),
-        eps=args_bounds_check(args, 7, replacement=1e-05),
-        cudnn_enabled=args_bounds_check(args, 8, replacement=False),
+        weight=args[1],
+        bias=args[2],
+        running_mean=args[3],
+        running_var=args[4],
+        training=args[5],
+        momentum=args[6],
+        eps=args[7],
+        cudnn_enabled=args_bounds_check(args, 8, replacement=True),
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
 def aten_ops_layer_norm(
     network: TRTNetwork,
@@ -58,9 +60,9 @@ def aten_ops_layer_norm(
         name,
         input=args[0],
         normalized_shape=args[1],
-        weight=args_bounds_check(args, 2, replacement=1),
-        bias=args_bounds_check(args, 3, replacement=0),
-        eps=args_bounds_check(args, 4, replacement=1e-05),
+        weight=args[2],
+        bias=args[3],
+        eps=args[4],
         cudnn_enable=args_bounds_check(args, 5, replacement=True),
     )
 

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

@@ -29,13 +29,13 @@ def batch_norm(
     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[float],
+    weight: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    running_mean: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    running_var: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    training: bool,
+    momentum: float,
+    eps: float,
     cudnn_enabled: bool,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, TRTTensor):
@@ -47,8 +47,20 @@ def batch_norm(
     if has_dynamic_shape(input.shape):
         assert input.shape[1] != -1, "Channel dim can't be dynamic for batch norm."
 
+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)
+
+    if running_mean is None:
+        running_mean = np.array(0.0)
+
+    if running_var is None:
+        running_var = np.array(1.0)
+
     scale = cast(torch.Tensor, to_numpy(weight)) / np.sqrt(
-        cast(torch.Tensor, to_numpy(running_var)) + cast(float, eps)
+        cast(torch.Tensor, to_numpy(running_var)) + eps
     )
 
     bias = to_numpy(bias) - to_numpy(running_mean) * scale
@@ -91,9 +103,9 @@ def layer_norm(
     name: str,
     input: TRTTensor,
     normalized_shape: List[int],
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    eps: List[float],
+    weight: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    eps: float,
     cudnn_enable: bool,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, trt.tensorrt.ITensor):
@@ -102,6 +114,12 @@ def layer_norm(
             "of the TensorRT region!"
         )
 
+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)
+
     gamma = (
         weight.detach().cpu().float().numpy()
         if isinstance(weight, torch.Tensor)
@@ -152,16 +170,22 @@ def layer_norm_no_plugin(
     name: str,
     input: TRTTensor,
     normalized_shape: List[int],
-    weight: torch.Tensor,
-    bias: torch.Tensor,
-    eps: List[float],
+    weight: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
+    eps: float,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
             f"LayerNorm received input {input} that is not part "
             "of the TensorRT region!"
         )
 
+    if weight is None:
+        weight = np.array(1.0)
+
+    if bias is None:
+        bias = np.array(0.0)
+
     shape = weight.shape
     broadcasted_shape = (1,) * (len(input.shape) - len(shape)) + shape
     gamma = to_numpy(weight.reshape(*shape))