[tvm] branch unity updated: [Unity][Frontend] from_fx keeps parameters in order (#14214)

[tq...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/unity by this push:
     new 50992d662f [Unity][Frontend] from_fx keeps parameters in order (#14214)
50992d662f is described below

commit 50992d662f357fbf9b43edc4471f5ecd216568ab
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Mon Mar 6 16:46:21 2023 -0500

    [Unity][Frontend] from_fx keeps parameters in order (#14214)
    
    Previously `from_fx` iterates model's `named_parameters()` to fetch all
    model's parameters. However, the order of iteration is not deterministic
    and might cause flakiness and instability in deployment.
    
    This PR removes the undeterminism by using `sorted(named_parameters())`.
    In this way we are guaranteed to have a deterministic list of function
    parameters.
---
 python/tvm/relax/__init__.py                     |  1 +
 python/tvm/relax/frontend/common.py              |  2 +-
 python/tvm/relax/frontend/torch/fx_translator.py |  2 +-
 tests/python/relax/test_frontend_from_fx.py      | 12 ++++++------
 4 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index e86f8c6074..bbd2040dd9 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -88,6 +88,7 @@ from . import block_builder
 from . import op
 from . import struct_info
 from . import backend
+from . import frontend
 
 # VM
 from .vm_build import build, Executable
diff --git a/python/tvm/relax/frontend/common.py b/python/tvm/relax/frontend/common.py
index e4432d8c67..9904324df4 100644
--- a/python/tvm/relax/frontend/common.py
+++ b/python/tvm/relax/frontend/common.py
@@ -42,7 +42,7 @@ def detach_params(mod: tvm.IRModule) -> Tuple[tvm.IRModule, Dict[str, List[tvm.n
     detached_mod = tvm.IRModule()
     params_dict = dict()
     for gv, func in mod.functions.items():
-        if "params" in func.attrs:
+        if func.attrs is not None and "params" in func.attrs:
             params = list(func.attrs["params"])
             if not all([isinstance(param, tvm.nd.NDArray) for param in params]):
                 raise ValueError(
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py
index d865984a14..a157b80337 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -923,7 +923,7 @@ class TorchFXImporter:
         params = []
         if keep_params_as_input:
             func_attrs = {"num_input": len(inputs)}
-            for name, param in model.named_parameters():
+            for name, param in sorted(model.named_parameters(), key=lambda x: x[0]):
                 shape = param.data.shape
                 dtype = self._convert_data_type(str(param.data.dtype))
                 inputs.append(relax.Var(name, relax.TensorStructInfo(shape, dtype)))
diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py
index 458efd7fcb..3447c99e5c 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2189,15 +2189,15 @@ def test_keep_params():
         @R.function
         def main(
             input_1: R.Tensor((1, 3, 10, 10), dtype="float32"),
-            w1: R.Tensor((6, 3, 7, 7), dtype="float32"),
-            w2: R.Tensor((6,), dtype="float32"),
+            w1: R.Tensor((6,), dtype="float32"),
+            w2: R.Tensor((6, 3, 7, 7), dtype="float32"),
         ) -> R.Tensor((1, 6, 4, 4), dtype="float32"):
             R.func_attr({"num_input": 1})
             # block 0
             with R.dataflow():
                 lv1: R.Tensor((1, 6, 4, 4), dtype="float32") = R.nn.conv2d(
                     input_1,
-                    w1,
+                    w2,
                     strides=[1, 1],
                     padding=[0, 0, 0, 0],
                     dilation=[1, 1],
@@ -2206,7 +2206,7 @@ def test_keep_params():
                     out_layout="NCHW",
                     out_dtype="float32",
                 )
-                lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w2, [1, 6, 1, 1])
+                lv2: R.Tensor((1, 6, 1, 1), dtype="float32") = R.reshape(w1, [1, 6, 1, 1])
                 lv3: R.Tensor((1, 6, 4, 4), dtype="float32") = R.add(lv1, lv2)
                 gv: R.Tensor((1, 6, 4, 4), dtype="float32") = lv3
                 R.output(gv)
@@ -2225,8 +2225,8 @@ def test_keep_params():
         assert tuple(x.value for x in param_var.struct_info.shape.values) == param_ndarray.shape
         assert param_var.struct_info.dtype == param_ndarray.dtype
 
-    tvm.testing.assert_allclose(params[0].numpy(), model.conv.weight.detach().numpy())
-    tvm.testing.assert_allclose(params[1].numpy(), model.conv.bias.detach().numpy())
+    tvm.testing.assert_allclose(params[0].numpy(), model.conv.bias.detach().numpy())
+    tvm.testing.assert_allclose(params[1].numpy(), model.conv.weight.detach().numpy())
 
 
 @tvm.testing.requires_gpu