[jit] Add weak script modules by driazati · Pull Request #12682 · pytorch/pytorch

driazati · 2018-10-15T22:03:14Z

Adds support for weak script modules created that get compiled to ScriptModules once added as a submodule of a ScriptModule:

@weak_module
class Test(torch.nn.Module):
	...
	@weak_script_method
	def forward(self, x):
		...

zdevito

Good progress! I added a bunch of comments. Largest thing I see is the need for careful tests that will hit all the corner cases that can happen.

Sign in to view

…_mod

zdevito

I am confused by the logic for maintaining information about weak_module classes and instances. I don't see why we need WeakModuleInstance to exist and WeakModule is awkward as a named tuple, considering that it is being mutated by reallocating it. It would help to organize this functionality into a clear API of classes and functions. As one example, consider: weak_modules.get(value.__class__). This line appears in setattr of the proxy class, which has multiple responsibilities. But what it is really saying is is_weak_script_annotated(value.class). Even though that function is a one liner, it is a lot easier to understand the code if it is refactored this way. Similarly for the rest of the logic here: it helps to decompose things into small functions with meaningful names and to keep the actual data structures (in this case, the weak dict tables), hidden behind this API. That way the implementation details do not leak into client code.

Sign in to view

+                return x + self.weak_submodule(x) + self.strong_submodule(x)
+
+        class Strong(torch.jit.ScriptModule):
+            __constants__ = ['constant']


Sign in to view

+
+            @torch.jit.weak_script_method
+            def forward(self, x):
+                return x * x


Sign in to view

        Symbol::fromQualString(py::str(builtin_name)), c10::nullopt);
  }

+  auto compiled_mod =


Sign in to view

+
+
+def _try_get_weak_module(mod):
+    weak_mods_of_class = weak_modules.get(mod.__class__)


Sign in to view

+
+def _try_get_weak_module(mod):
+    weak_mods_of_class = weak_modules.get(mod.__class__)
+    if weak_mods_of_class is not None:


Sign in to view

+    return cls
+
+
+def _get_weak_stubs(mod):


Sign in to view

 BatchTensor = torch._C._jit.BatchTensor
 compiled_weak_fns = weakref.WeakKeyDictionary()
+weak_script_methods = weakref.WeakKeyDictionary()
+weak_modules = weakref.WeakKeyDictionary()


Sign in to view

+def _try_get_weak_module(mod):
+    weak_mods_of_class = weak_modules.get(mod.__class__)
+    if weak_mods_of_class is not None:
+        weak_mod = weak_mods_of_class.get(mod)


Sign in to view

+    if weak_module.instances is None:
+        # Generate stubs and add to entry
+        stubs = _get_weak_stubs(mod)
+        weak_modules[cls] = WeakModule(weakref.WeakKeyDictionary(), stubs)


Sign in to view

+
+
+def _try_get_weak_module(mod):
+    weak_mods_of_class = weak_modules.get(mod.__class__)


…_mod

zdevito

This looks much better! I had one minor thing about caching results in _make_strong. Other than that I think this is ready to go!

Sign in to view

+    """
+    Converts a weak module into a subclass of ScriptModule
+    """
+    stubs = _weak_types.get(type(mod))["method_stubs"]


facebook-github-bot

driazati has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

driazati has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

driazati has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

driazati has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Depends on #12682 ([stacked diff](https://github.com/driazati/pytorch/compare/weak_mod...driazati:mod_conv1)) * Adds tests for weak module conversion that creates a `ScriptModule` that uses the weak module and checks its graph * Adds `torch._jit_internal.weak_module` tags to modules that already work * `Sigmoid` * `Tanh` * `Hardshrink` * `PReLU` * `Softsign` * `Tanhshrink` * `PairwiseDistance` Pull Request resolved: #12966 Differential Revision: D10559557 Pulled By: driazati fbshipit-source-id: dc4bea3aa744b3c44d4fa7dceefd97e951f824d0

Summary: Adds support for weak script modules created that get compiled to `ScriptModule`s once added as a submodule of a `ScriptModule`: ```python weak_module class Test(torch.nn.Module): ... weak_script_method def forward(self, x): ... ``` Pull Request resolved: pytorch#12682 Differential Revision: D10458626 Pulled By: driazati fbshipit-source-id: 10ae23cb83cdafc4646cee58f399e14b2e60acd4

Summary: Depends on pytorch#12682 ([stacked diff](https://github.com/driazati/pytorch/compare/weak_mod...driazati:mod_conv1)) * Adds tests for weak module conversion that creates a `ScriptModule` that uses the weak module and checks its graph * Adds `torch._jit_internal.weak_module` tags to modules that already work * `Sigmoid` * `Tanh` * `Hardshrink` * `PReLU` * `Softsign` * `Tanhshrink` * `PairwiseDistance` Pull Request resolved: pytorch#12966 Differential Revision: D10559557 Pulled By: driazati fbshipit-source-id: dc4bea3aa744b3c44d4fa7dceefd97e951f824d0

zdevito reviewed Oct 16, 2018

View reviewed changes

David Riazati added 11 commits October 16, 2018 10:12

Weak script module

942df07

Constants

b430a4c

Add tests

8338a5f

Cleanup

989cb77

Address comments

623bad1

Nested test

537b641

Buffers

d2e2f9e

Merge branch 'master' of https://github.com/pytorch/pytorch into weak…

ed7156f

…_mod

Map holds weak scripts instead of the original object

23506b0

Dont re-parse AST

052b7f4

Cleanup

1ded7bc

zdevito reviewed Oct 17, 2018

View reviewed changes

David Riazati added 5 commits October 17, 2018 16:54

Simplify/document module tracking logic, add tests

2f0a36a

Lint fix

ce2035e

Fix PY2/3 error

1b7c5b4

Merge branch 'master' of https://github.com/pytorch/pytorch into weak…

548494d

…_mod

Clean up names

a89cd55

zdevito approved these changes Oct 18, 2018

View reviewed changes

Comment thread torch/jit/__init__.py

"""

Converts a weak module into a subclass of ScriptModule

"""

stubs = _weak_types.get(type(mod))["method_stubs"]

This comment was marked as off-topic.

Sign in to view

Cache proxy

a305e39

facebook-github-bot reviewed Oct 19, 2018

View reviewed changes

David Riazati added 2 commits October 22, 2018 09:02

Merge

e761810

Move weak module annotations out to _jit_internal

8901777

driazati closed this Oct 22, 2018

driazati reopened this Oct 22, 2018

facebook-github-bot reviewed Oct 22, 2018

View reviewed changes

driazati closed this Oct 22, 2018

driazati reopened this Oct 22, 2018

driazati closed this Oct 22, 2018

driazati reopened this Oct 22, 2018

Style

207a01b

facebook-github-bot reviewed Oct 22, 2018

View reviewed changes

Lint fix

e91decf

facebook-github-bot reviewed Oct 22, 2018

View reviewed changes

driazati mentioned this pull request Oct 23, 2018

[jit] Make 7 nn modules into weak modules #12966

Closed

facebook-github-bot closed this in af78d4c Oct 23, 2018

ezyang added the merged label Jun 25, 2019



		def _try_get_weak_module(mod):
		weak_mods_of_class = weak_modules.get(mod.__class__)

Conversation

driazati commented Oct 15, 2018

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as off-topic.

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants