[jit] Support recursive ModuleList / Sequential

davidriazati · davidriazati · commit 67eeeed94d8c · 2019-06-03T15:02:11.000-07:00
Adds support for recursively compiling `nn.Sequential` and
`nn.ModuleList`. When either is used, it is converted to a
`jit._ConstModuleList` or `jit._ConstSequential` as necessary. Due to
this, we don't need to add it to `__constants__` since it's made
constant on demand.

This PR also moves the recursive script tests out to their own class
`TestRecursiveScript` (the added test is called `test_iterable_modules`)
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -13133,86 +13133,6 @@ def is_tensor_value(item):
                     continue
                 self.assertEqual(item[1], loaded_item)
 
-    def test_script_recurse(self):
-        def a_python_fn(a, b, c):
-            return a + b + c
-
-        with torch.jit._enable_recursive_script():
-            @torch.jit.script
-            def a_script_fn(d, e, f):
-                return a_python_fn(d, e, f)
-
-        graph = str(a_script_fn.graph)
-        FileCheck().check("aten::add").run(graph)
-        FileCheck().check_not("a_python_fn").run(graph)
-        t = torch.ones(2, 2)
-        self.assertEqual(a_script_fn(t, t, t), t + t + t)
-
-    def test_module_recursive(self):
-        class Other(torch.nn.Module):
-            __constants__ = ['x']
-
-            def __init__(self, x):
-                super(Other, self).__init__()
-                self.x = x
-                self.param = torch.nn.Parameter(torch.ones(2, 2))
-
-            def some_unscriptable_method(self):
-                a = 2
-                a = [2]
-                return a
-
-            def forward(self, t):
-                return t + self.x + self.param
-
-
-        class M(torch.nn.Module):
-            __constants__ = ['x']
-
-            def __init__(self):
-                super(M, self).__init__()
-                self.other = Other(200)
-
-            def forward(self, t):
-                return self.other(t) * 2
-
-        with torch.jit._enable_recursive_script():
-            sm = torch.jit.script(M())
-
-        self.assertExportImportModule(sm, (torch.ones(2, 2),))
-
-    def test_module_function_export(self):
-        class Other(torch.nn.Module):
-            __constants__ = ['x']
-
-            def __init__(self, x):
-                super(Other, self).__init__()
-                self.x = x
-                self.param = torch.nn.Parameter(torch.ones(2, 2))
-
-            @torch.jit.export
-            def some_entry_point(self, y):
-                return y + 20
-
-            def forward(self, t):
-                return t + self.x + self.param
-
-
-        class M(torch.nn.Module):
-            __constants__ = ['x']
-
-            def __init__(self):
-                super(M, self).__init__()
-                self.other = Other(200)
-
-            def forward(self, t):
-                return self.other(t) * 2
-
-        with torch.jit._enable_recursive_script():
-            sm = torch.jit.script(M())
-
-        self.assertExportImportModule(sm, (torch.ones(2, 2),))
-
     @unittest.skipIf(IS_WINDOWS or IS_SANDCASTLE, "NYI: TemporaryFileName support for Windows or Sandcastle")
     def test_old_models_bc(self):
         model = {
@@ -13475,6 +13395,125 @@ def foo(a):
             foo(1)
 
 
+class TestRecursiveScript(JitTestCase):
+    """
+    Tests in this class are all run under `with torch.jit._enable_recursive_script()`
+    """
+    def run(self, result=None):
+        with torch.jit._enable_recursive_script():
+            super(TestRecursiveScript, self).run(result)
+
+    def checkModule(self, nn_module, args):
+        """
+        Check that a nn.Module's results in Script mode match eager and that it
+        can be exported
+        """
+        sm = torch.jit.script(nn_module)
+
+        eager_out = nn_module(*args)
+        script_out = sm(*args)
+
+        self.assertEqual(eager_out, script_out)
+        self.assertExportImportModule(sm, args)
+
+        return sm
+
+    def test_script_basic(self):
+        def a_python_fn(a, b, c):
+            return a + b + c
+
+        @torch.jit.script
+        def a_script_fn(d, e, f):
+            return a_python_fn(d, e, f)
+
+        graph = str(a_script_fn.graph)
+        FileCheck().check("aten::add").run(graph)
+        FileCheck().check_not("a_python_fn").run(graph)
+        t = torch.ones(2, 2)
+        self.assertEqual(a_script_fn(t, t, t), t + t + t)
+
+    def test_module_basic(self):
+        class Other(torch.nn.Module):
+            __constants__ = ['x']
+
+            def __init__(self, x):
+                super(Other, self).__init__()
+                self.x = x
+                self.param = torch.nn.Parameter(torch.ones(2, 2))
+
+            def some_unscriptable_method(self):
+                a = 2
+                a = [2]
+                return a
+
+            def forward(self, t):
+                return t + self.x + self.param
+
+
+        class M(torch.nn.Module):
+            __constants__ = ['x']
+
+            def __init__(self):
+                super(M, self).__init__()
+                self.other = Other(200)
+
+            def forward(self, t):
+                return self.other(t) * 2
+
+        sm = torch.jit.script(M())
+
+        self.assertExportImportModule(sm, (torch.ones(2, 2),))
+
+    def test_module_function_export(self):
+        class Other(torch.nn.Module):
+            __constants__ = ['x']
+
+            def __init__(self, x):
+                super(Other, self).__init__()
+                self.x = x
+                self.param = torch.nn.Parameter(torch.ones(2, 2))
+
+            @torch.jit.export
+            def some_entry_point(self, y):
+                return y + 20
+
+            def forward(self, t):
+                return t + self.x + self.param
+
+
+        class M(torch.nn.Module):
+            __constants__ = ['x']
+
+            def __init__(self):
+                super(M, self).__init__()
+                self.other = Other(200)
+
+            def forward(self, t):
+                return self.other(t) * 2
+
+        sm = torch.jit.script(M())
+        self.assertExportImportModule(sm, (torch.ones(2, 2),))
+
+    def test_iterable_modules(self):
+        class M(torch.nn.Module):
+            def __init__(self):
+                super(M, self).__init__()
+                self.sequential = nn.Sequential(
+                    nn.Linear(5, 5),
+                    nn.Linear(5, 5),
+                    nn.Sequential(nn.Linear(5, 5), nn.Linear(5, 5))
+                )
+                self.module_list = nn.ModuleList([nn.Linear(5, 5), nn.Linear(5, 5)])
+
+            def forward(self, x):
+                for mod in self.module_list:
+                    x += mod(x)
+                x += self.sequential(x)
+                return x
+
+        self.checkModule(M(), (torch.randn(5, 5),))
+
+
 class MnistNet(nn.Module):
     def __init__(self):
         super(MnistNet, self).__init__()
diff --git a/torch/csrc/jit/script/python_sugared_value.cpp b/torch/csrc/jit/script/python_sugared_value.cpp
@@ -275,7 +275,6 @@ std::shared_ptr<SugaredValue> ModuleValue::attr(
 
   // This can also be a call to a non-script module, or a plain
   // python method. If so return this as a python value.
-
   py::object overloads =
       py_module_.attr("_overloads").attr("get")(field, py::none());
   if (!overloads.is_none()) {
@@ -327,7 +326,7 @@ std::shared_ptr<SugaredValue> ModuleValue::attr(
         module_->register_module(field, submodule);
         auto v = module_->find_module(field);
         return std::make_shared<ModuleValue>(
-            m.graph()->insertGetAttr(self_, field), v, attr);
+            m.graph()->insertGetAttr(self_, field), v, result);
       }
     } else if (py::isinstance<py::function>(attr)) {
       auto stub = py::module::import("torch.jit")
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -940,11 +940,26 @@ def env(key):
 
     return env
 
+def _create_constant_iterable_module(module):
+    for i, submodule in enumerate(module):
+        if isinstance(submodule, (ModuleList, Sequential)):
+            # Make each item in the module a constant
+            module[i] = _create_constant_iterable_module(submodule)
+
+    if isinstance(module, Sequential):
+        return _ConstSequential(module)
+    elif isinstance(module, ModuleList):
+        return _ConstModuleList(module)
+    else:
+        raise RuntimeError("Only nn.ModuleList and nn.Sequential can be made "
+                           "into constant modules, found {}".format(module))
+
 
 def _make_strong_submodule(field, module, parent):
     if field not in parent._modules:
         # It's not a submodule, don't do anything
         return None
+
     return _convert_to_script_module(module)
 
 
@@ -953,6 +968,9 @@ def _try_compile_fn(fn):
         # Don't do anything for @ignore'd functions
         return None
 
+    if isinstance(fn, torch.nn.Module):
+        return None
+
     # We don't have the actual scope where the function was defined, but we can
     # extract the necessary info from the closed over variables on the function
     # object
@@ -1237,8 +1255,6 @@ def _create_methods_from_stubs(self, stubs):
 # has run. This has to occur after the user-defined __init__ so that
 # submodules and parameters are initialized _before_ the script compiler
 # resolve references to `self.param` or `self.module`.
-
-
 class ScriptMeta(type):
     # this has to inherit from pybind11's metaclass otherwise we get
     # issues because ScriptModule inherits from torch._C.ScriptModule,
@@ -1260,6 +1276,7 @@ def __init__(cls, name, bases, attrs):
                 cls._methods[v.original_method.__name__] = v
 
         original_init = getattr(cls, '__init__', lambda self: None)
+        print("set overloads", cls, name)
         cls._overloads = dict(getattr(cls, '__overloads__', {}))
 
         # after the user's __init__ register all the script methods
@@ -1537,6 +1554,11 @@ def graph_for(self, *args, **kwargs):
             return self.forward.graph_for(*args, **kwargs)
 
     class WeakScriptModuleProxy(ScriptModule):
+        # TODO: [weak script refactor]
+        # WeakScriptModule proxy should be deleted since its functionality is
+        # subsumed by recursive scripting, and the copying code in init moved
+        # to a function to create a ScriptModule from an nn.Module without
+        # making a WeakScriptModuleProxy
         """
         Copies the parameters, buffers, constants, attributes, and submodules
         of an nn.Module into itself.
@@ -1568,7 +1590,13 @@ def __init__(self, original, stubs):
                 elif isinstance(item, (Parameter, Module, Attribute)):
                     if isinstance(item, (ModuleList, Sequential)):
                         # These are in __constants__, so ignore them here
-                        continue
+
+                        if not torch._C._jit_recursive_script():
+                            # For recursive script, these are constantified after
+                            # they are used, so they don't need to be in constants.
+                            # The `continue` here should be deleted along with
+                            # [weak script refactor]
+                            continue
                     ScriptModule.__setattr__(self, name, item)
 
             # Copy buffers
@@ -1651,6 +1679,10 @@ def _convert_to_script_module(mod, methods=None):
     `('forward',)`. Methods accessed in forward are scripted on demand if
     `_enable_recursive_script()` is used.
     """
+    if isinstance(mod, (ModuleList, Sequential)):
+        # Create constant versions for the iterable modules
+        return _create_constant_iterable_module(mod)
+
     if methods is None:
         methods = ('forward',)
     exported = []
