[acc_tracer] Add test coverage for retracing (#71752)

jfix71 · pytorchmergebot · commit 66939e3b94fa · 2022-01-26T04:30:18.000Z
Summary: Pull Request resolved: #71752 Added coverage for reshape specifically which required a fix. The problem for `acc_ops.reshape` as best as I understand: - `torch.reshape` requires the `shape` arg to be a `tuple` of `ints` - If `torch.reshape` is passed a `tuple` where the first element is not an `int`, it throws a TypeError e.g. `TypeError: reshape(): argument 'shape' (position 2) must be tuple of ints, not tuple` - If the `shape` we're reshaping to is an FX Proxy then this type error will be thrown. This happens when the first element of the `shape` tuple is a Proxy because it's input-dependent. - As a workaround we use `tensor.reshape` instead of `torch.reshape`, which doesn't do equivalent type checking for a `tuple` of `ints`. Also remove unnecessary `acc_utils.get_field_from_acc_out_ty()` with cast to `TensorMetadata`. Test Plan: Added test coverage Reviewed By: yinghai Differential Revision: D33760455 fbshipit-source-id: bff5563bf9e3d9e9318901b56211151d2c0e4eb2 (cherry picked from commit d5c1b97)
diff --git a/test/fx_acc/test_acc_tracer.py b/test/fx_acc/test_acc_tracer.py
@@ -37,6 +37,11 @@ def _make_model_unit_test(
             torch.testing.assert_allclose(model(input), traced(input))
         else:
             self.assertTrue(torch.equal(model(input), traced(input)))
+        traced_again = acc_tracer.trace(traced, [input])
+        if enable_allclose:
+            torch.testing.assert_allclose(model(input), traced_again(input))
+        else:
+            self.assertTrue(torch.equal(model(input), traced_again(input)))
 
     def _make_acc_op_function_test(
         self,
@@ -89,30 +94,47 @@ def forward(self, a: torch.Tensor) -> torch.Tensor:
 
         ref_outputs = m(a)
         outputs = traced(a)
+        traced_again = acc_tracer.trace(m, [a])
+        outputs_again = traced_again(a)
         if isinstance(ref_outputs, torch.Tensor):
             ref_outputs = [ref_outputs]
             outputs = [outputs]
+            outputs_again = [outputs_again]
 
-        for ref_output, output in zip(ref_outputs, outputs):
+        for ref_output, output, output_again in zip(
+            ref_outputs, outputs, outputs_again
+        ):
             if enable_allclose:
                 torch.testing.assert_allclose(
                     torch.nan_to_num(ref_output), torch.nan_to_num(output)
                 )
+                torch.testing.assert_allclose(
+                    torch.nan_to_num(ref_output), torch.nan_to_num(output_again)
+                )
             else:
                 self.assertTrue(
                     torch.equal(torch.nan_to_num(ref_output), torch.nan_to_num(output))
                 )
+                self.assertTrue(
+                    torch.equal(
+                        torch.nan_to_num(ref_output), torch.nan_to_num(output_again)
+                    )
+                )
 
     def test_sum(self):
         self._make_acc_op_function_test(acc_ops.sum, torch.sum)
         self._make_acc_op_function_test(acc_ops.sum, torch.sum, dim=(1,), keepdim=True)
 
     def test_mean(self):
         self._make_acc_op_function_test(acc_ops.mean, torch.mean)
-        self._make_acc_op_function_test(acc_ops.mean, torch.mean, dim=(1,), keepdim=True)
+        self._make_acc_op_function_test(
+            acc_ops.mean, torch.mean, dim=(1,), keepdim=True
+        )
 
     def test_pad(self):
-        self._make_acc_op_function_test(acc_ops.pad, torch.nn.functional.pad, pad=(2, 0))
+        self._make_acc_op_function_test(
+            acc_ops.pad, torch.nn.functional.pad, pad=(2, 0)
+        )
 
     def test_max(self):
         def torch_max(x, *args, **kwargs):
@@ -504,7 +526,9 @@ def forward(self, a: torch.Tensor) -> torch.Tensor:
                 self.assertEqual(node.kwargs["running_mean"], bn_mean)
                 self.assertEqual(node.kwargs["running_var"], bn_var)
                 self.assertEqual(node.kwargs["acc_out_ty"][6]["scale"], bn_scale)
-                self.assertEqual(node.kwargs["acc_out_ty"][6]["zero_point"], bn_zero_point)
+                self.assertEqual(
+                    node.kwargs["acc_out_ty"][6]["zero_point"], bn_zero_point
+                )
                 bn = node
             elif node.op == "output":
                 self.assertEqual(bn, node.args[0])
@@ -1230,7 +1254,9 @@ def test_dropout(self):
     def test_stochastic_depth(self):
         self._make_acc_op_function_test(
             None,
-            lambda x, p, mode, training: torchvision.ops.stochastic_depth(x, p=p, mode=mode, training=training),
+            lambda x, p, mode, training: torchvision.ops.stochastic_depth(
+                x, p=p, mode=mode, training=training
+            ),
             input_shape=(1, 2, 3),
             p=0.5,
             mode="row",
@@ -1387,7 +1413,9 @@ def test_relu(self):
         self._make_acc_op_function_test(acc_ops.relu, torch.relu)
 
     def test_leaky_relu(self):
-        self._make_acc_op_function_test(acc_ops.leaky_relu, torch.nn.functional.leaky_relu)
+        self._make_acc_op_function_test(
+            acc_ops.leaky_relu, torch.nn.functional.leaky_relu
+        )
 
     def test_elu(self):
         self._make_acc_op_function_test(acc_ops.elu, torch.nn.functional.elu)
@@ -1472,11 +1500,17 @@ def test_div(self):
         self._make_acc_op_function_test(acc_ops.div, lambda x: x / 2)
 
     def test_floor_div(self):
-        self._make_acc_op_function_test(acc_ops.floor_div, lambda x: torch.div(x, 2, rounding_mode="floor"))
+        self._make_acc_op_function_test(
+            acc_ops.floor_div, lambda x: torch.div(x, 2, rounding_mode="floor")
+        )
 
     def test_trunc_div(self):
-        self._make_acc_op_function_test(acc_ops.trunc_div, lambda x: torch.div(x, 2, rounding_mode="trunc"))
-        self._make_acc_op_function_test(acc_ops.trunc_div, lambda x: torch.floor_divide(x, 2))
+        self._make_acc_op_function_test(
+            acc_ops.trunc_div, lambda x: torch.div(x, 2, rounding_mode="trunc")
+        )
+        self._make_acc_op_function_test(
+            acc_ops.trunc_div, lambda x: torch.floor_divide(x, 2)
+        )
 
     def test_view(self):
         """
@@ -1907,6 +1941,22 @@ def test_cumsum(self):
     def test_chunk(self):
         self._make_acc_op_function_test(acc_ops.chunk, torch.chunk, chunks=2, dim=0)
 
+    def test_retrace_reshape(self):
+        """
+        Retrace reshape to verify it's retraceable.
+        """
+
+        class TestModule(torch.nn.Module):
+            def forward(self, a: torch.Tensor) -> torch.Tensor:
+                return a.reshape(a.size()[0], 1, 2)
+
+        m = TestModule()
+        a = torch.randn(2, 2)
+        gm = acc_tracer.trace(m, [a])
+        self.assertTrue(torch.equal(m(a), gm(a)))
+        gm_retrace = acc_tracer.trace(gm, [a])
+        self.assertTrue(torch.equal(m(a), gm_retrace(a)))
+
     def test_all_acc_ops_registered(self):
         self.assertEqual(
             acc_normalizer._acc_ops,
diff --git a/torch/fx/experimental/fx2trt/converters/acc_ops_converters.py b/torch/fx/experimental/fx2trt/converters/acc_ops_converters.py
@@ -7,7 +7,6 @@
 import tensorrt as trt
 import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
-import torch.fx.experimental.fx_acc.acc_utils as acc_utils
 from torch.fx.experimental.fx2trt.converter_registry import tensorrt_converter
 from torch.fx.experimental.fx2trt.types import *  # noqa: F403
 from torch.fx.experimental.fx2trt.utils import (
@@ -16,6 +15,7 @@
 )
 from torch.fx.immutable_collections import immutable_list
 from torch.fx.node import Target, Argument
+from torch.fx.passes.shape_prop import TensorMetadata
 
 from .converter_utils import *  # noqa: F403
 
@@ -1376,7 +1376,7 @@ def acc_ops_reshape(
             "of the TensorRT region!"
         )
 
-    shape = acc_utils.get_field_from_acc_out_ty(kwargs["acc_out_ty"], "shape")  # type: ignore[arg-type]
+    shape = TensorMetadata(*kwargs["acc_out_ty"]).shape  # type: ignore[misc]
     if network.has_implicit_batch_dimension:
         shape = shape[1:]
 
@@ -1887,10 +1887,10 @@ def acc_ops_quantize_per_tensor(
         raise RuntimeError(f"{name} received input {input_val} that is not part "
                            "of the TensorRT region!")
 
-    qparams = acc_utils.get_field_from_acc_out_ty(kwargs["acc_out_ty"], "qparams")  # type: ignore[arg-type]
+    qparams = TensorMetadata(*kwargs["acc_out_ty"]).qparams  # type: ignore[misc]
     q_scale = qparams["scale"]
     q_zero_point = qparams["zero_point"]
-    dtype = acc_utils.get_field_from_acc_out_ty(kwargs["acc_out_ty"], "dtype")  # type: ignore[arg-type]
+    dtype = TensorMetadata(*kwargs["acc_out_ty"]).dtype  # type: ignore[misc]
     if dtype not in (torch.quint8, torch.qint8, torch.qint32):
         raise RuntimeError("Only support (torch.quint8, torch.qint8, torch.qint32) "
                            f"quantized type in quantize_per_tensor, get {dtype}.")
@@ -1923,11 +1923,11 @@ def acc_ops_quantize_per_channel(
         raise RuntimeError(f"{name} received input {input_val} that is not part "
                            "of the TensorRT region!")
 
-    qparams = acc_utils.get_field_from_acc_out_ty(kwargs["acc_out_ty"], "qparams")  # type: ignore[arg-type]
+    qparams = TensorMetadata(*kwargs["acc_out_ty"]).qparams  # type: ignore[misc]
     q_per_channel_scales = qparams["scale"]
     q_per_channel_zero_points = qparams["zero_point"]
     q_per_channel_axis = qparams["axis"]
-    dtype = acc_utils.get_field_from_acc_out_ty(kwargs["acc_out_ty"], "dtype")  # type: ignore[arg-type]
+    dtype = TensorMetadata(*kwargs["acc_out_ty"]).dtype  # type: ignore[misc]
     if dtype not in (torch.quint8, torch.qint8, torch.qint32):
         raise RuntimeError("Only support (torch.quint8, torch.qint8, torch.qint32) "
                            f"quantized type in quantize_per_tensor, get {dtype}.")
@@ -1970,7 +1970,7 @@ def acc_ops_dequantize(
         raise RuntimeError(f"{name} received input {input_val} that is not part "
                            "of the TensorRT region!")
 
-    qparams = acc_utils.get_field_from_acc_out_ty(input_val_tensor_meta, "qparams")  # type: ignore[arg-type]
+    qparams = TensorMetadata(*input_val_tensor_meta).qparams  # type: ignore[misc]
     qscheme = qparams["qscheme"]
     if qscheme == torch.per_tensor_affine:
         q_scale = qparams["scale"]
@@ -1990,7 +1990,7 @@ def acc_ops_dequantize(
     else:
         raise RuntimeError("Unsupported qscheme in dequantize: {qscheme}")
 
-    dtype = acc_utils.get_field_from_acc_out_ty(input_val_tensor_meta, "dtype")  # type: ignore[arg-type]
+    dtype = TensorMetadata(*input_val_tensor_meta).dtype  # type: ignore[misc]
 
     if dtype not in (torch.quint8, torch.qint8, torch.qint32):
         raise RuntimeError("Only support (torch.quint8, torch.qint8, torch.qint32) "
diff --git a/torch/fx/experimental/fx_acc/acc_ops.py b/torch/fx/experimental/fx_acc/acc_ops.py
@@ -16,7 +16,7 @@
     AccOpProperty,
     register_acc_op_properties,
 )
-from torch.fx.passes.shape_prop import _extract_tensor_metadata
+from torch.fx.passes.shape_prop import _extract_tensor_metadata, TensorMetadata
 
 this_arg_is_optional = True
 move_to_qparams = True
@@ -33,7 +33,7 @@ def linear(*, input, weight, bias):
 @register_acc_op
 def quantized_linear(*, input, weight, bias, acc_out_ty=None):
     assert acc_out_ty is not None
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
+    qparams = TensorMetadata(*acc_out_ty).qparams
     return nn.quantized.functional.linear(
         input,
         weight,
@@ -462,11 +462,13 @@ def dropout_mapper(node: torch.fx.Node, mod: nn.Module):
     """
     return node.kwargs["input"]
 
+
 try:
     from torchvision.ops import stochastic_depth
 except Exception as e:
     warnings.warn(f"Unable to import torchvision related libraries.: {e}")
 else:
+
     @register_custom_acc_mapper_fn(
         op_and_target=("call_function", stochastic_depth),
         arg_replacement_tuples=[("input", "input")],
@@ -477,6 +479,7 @@ def stochastic_depth_mapper(node: torch.fx.Node, mod: nn.Module):
         """
         return node.kwargs["input"]
 
+
 @register_acc_op_properties(AccOpProperty.pointwise, AccOpProperty.unary)
 @register_acc_op_mapping(
     op_and_target=("call_function", nn.functional.hardtanh),
@@ -502,9 +505,7 @@ def hardsigmoid(*, input):
 def silu(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
     input_node = node.kwargs["input"]
     with node.graph.inserting_before(node):
-        sigmoid_node = node.graph.call_function(
-            sigmoid, kwargs={"input": input_node}
-        )
+        sigmoid_node = node.graph.call_function(sigmoid, kwargs={"input": input_node})
         sigmoid_node.meta = node.meta.copy()
         new_node = node.graph.call_function(
             mul, kwargs={"input": sigmoid_node, "other": input_node}
@@ -550,7 +551,7 @@ def hardswish_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
 @register_acc_op
 def quantized_add(*, input, other, acc_out_ty=None):
     assert acc_out_ty is not None
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
+    qparams = TensorMetadata(*acc_out_ty).qparams
     return torch.ops.quantized.add(
         input,
         other,
@@ -576,7 +577,7 @@ def quantized_add(*, input, other, acc_out_ty=None):
 @register_acc_op
 def quantized_mul(*, input, other, acc_out_ty=None):
     assert acc_out_ty is not None
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
+    qparams = TensorMetadata(*acc_out_ty).qparams
     return torch.ops.quantized.mul(
         input,
         other,
@@ -604,8 +605,8 @@ def quantized_mul(*, input, other, acc_out_ty=None):
 @register_acc_op
 def quantize_per_tensor(*, input, acc_out_ty=None):
     assert acc_out_ty is not None
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
-    dtype = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "dtype")
+    qparams = TensorMetadata(*acc_out_ty).qparams
+    dtype = TensorMetadata(*acc_out_ty).dtype
     return torch.quantize_per_tensor(
         input, qparams["scale"], qparams["zero_point"], dtype
     )
@@ -631,8 +632,8 @@ def quantize_per_tensor(*, input, acc_out_ty=None):
 @register_acc_op
 def quantize_per_channel(*, input, acc_out_ty=None):
     assert acc_out_ty is not None
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
-    dtype = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "dtype")
+    qparams = TensorMetadata(*acc_out_ty).qparams
+    dtype = TensorMetadata(*acc_out_ty).dtype
     return torch.quantize_per_channel(
         input,
         torch.tensor(qparams["scale"]),
@@ -1152,7 +1153,7 @@ def quantized_conv2d(
     padding_mode,
     acc_out_ty,
 ):
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
+    qparams = TensorMetadata(*acc_out_ty).qparams
     return torch.nn.quantized.functional.conv2d(
         input=input,
         weight=weight,
@@ -1359,7 +1360,7 @@ def tuple_construct(*, tensors):
 def quantized_batch_norm2d(
     *, input, running_mean, running_var, weight, bias, eps, acc_out_ty
 ):
-    qparams = acc_utils.get_field_from_acc_out_ty(acc_out_ty, "qparams")
+    qparams = TensorMetadata(*acc_out_ty).qparams
     return torch.ops.quantized.batch_norm2d(
         input,
         weight,
@@ -1573,9 +1574,7 @@ def custom_narrow_mapper(node: torch.fx.Node, mod: nn.Module) -> torch.fx.Node:
 @register_acc_op
 def reshape(*, input, acc_out_ty=None):
     assert acc_out_ty is not None
-    return torch.reshape(
-        input, tuple(acc_utils.get_field_from_acc_out_ty(acc_out_ty, "shape"))
-    )
+    return input.reshape(TensorMetadata(*acc_out_ty).shape)
 
 
 @register_custom_acc_mapper_fn(
@@ -1615,7 +1614,7 @@ def custom_tensor_reshape_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.
 @register_acc_op
 def to_dtype(input, acc_out_ty=None):
     assert acc_out_ty is not None
-    return input.to(dtype=acc_utils.get_field_from_acc_out_ty(acc_out_ty, "dtype"))
+    return input.to(dtype=TensorMetadata(*acc_out_ty).dtype)
 
 
 @register_custom_acc_mapper_fn(
diff --git a/torch/fx/experimental/fx_acc/acc_utils.py b/torch/fx/experimental/fx_acc/acc_utils.py
@@ -78,20 +78,6 @@ def is_acc_op_with_kwarg(
     return kwarg in inspect.signature(inspect.unwrap(target)).parameters
 
 
-def get_field_from_acc_out_ty(
-    acc_out_ty_or_dict: Union[Tuple, Dict[str, Any]], field: str
-):
-    """
-    After tracing NamedTuple inputs are converted to standard tuples, so we cannot
-    access them by name directly. Use this helper instead.
-    """
-    if isinstance(acc_out_ty_or_dict, dict):
-        acc_out_ty = acc_out_ty_or_dict["acc_out_ty"]
-    else:
-        acc_out_ty = acc_out_ty_or_dict
-    return acc_out_ty[TensorMetadata._fields.index(field)]
-
-
 def serialize_module_json_to_file(fx_module: GraphModule, fname: str):
     weights: Dict = {}
     serialized_json = json.dumps(serialize_module(fx_module, weights), indent=2)
