[jit] Add support for recursive compilation on Modules #20708

driazati · 2019-05-20T17:48:10Z

Following on #19747, this implements most of the torch.jit.script() changes laid out in #20939.

Still to do:

Accessing a method from Python does not add it as a ScriptMethod (so only @exported methods and forward are compiled)
Calling a method other than forward on a submodule doesn't work

Differential Revision: D15560490

suo · 2019-05-20T20:41:13Z

What would it take to get rid of __constants__ as well?

driazati · 2019-05-20T20:52:15Z

Similar to #19439 we could just make everything that we see in __init__ that matches the supported constant types a constant. This could be a bad UX though since there's no explicit thing saying "I won't change this". Alternatively everything could just be added as an attribute (unless explicitly made a constant in __constants__) and the UX would be preserved.

zdevito · 2019-05-21T17:27:40Z

This change deserves a design document before we land any code. There are some key questions which I am not clear on:

When something becomes strong, how is it related to the original module? Should it be the same object (I prefer this if possible)? Should we have coupled objects? If we have coupled objects what happens when we edit one but not the other?
How do we decide what methods should be scripted? Assuming just forward is used is not general, and making it configurable is not great because then it requires user input. We should make it possible that each method becomes scripted as required.

We should come up with a consistent design (across weak and strong modules) that we look at first.

zdevito

This is a good start. I have a few inline comments, but generally I think we should merge it and then discuss the next steps.

torch/_jit_internal.py

zdevito · 2019-05-29T05:23:55Z

torch/_jit_internal.py

+    saving. This decorator explicitly marks that a method should be included
+    even if it is not called from Python.
+    """
+    class_dict = get_class_attribute_dict(fn, frames_up=2)


I feel like this would be idiomatic if it just set a __torchscript_export property on fn and the class then searched its own __dict__ for things marked export. Searching the frames for what you hope to be the class dict is not great. What happens when we add @property and then we chain @export @property def. It would break.

torch/_jit_internal.py

zdevito · 2019-05-29T05:39:29Z

torch/jit/__init__.py



+def _make_strong_submodule(field, module, parent):
+    if field not in parent._modules:


how does this happen? There are nn.Modules in module that are not in the _modules list?

This was more of a sanity check though it's possible someone does self.__dict__["secret_module"] = nn.Linear()

zdevito · 2019-05-29T05:43:13Z

torch/jit/__init__.py

    class WeakScriptModuleProxy(ScriptModule):
        def __init__(self, original, stubs):
            # Guards behavior of __setattr__ and __getattr__ so ScriptModule
            # __init__ can run correctly


We need a plan to get rid of WeakScriptModuleProxy. Eventually all ScriptModules will just always be this class, and they will not longer be proxies at all. It is difficult to understand all the components we have after this patch, so we need some followup effort merging the concepts.

Most of the functionality here can operate directly on a ScriptModule instance instead of constructing a WeakScriptModuleProxy, so deleting it should be pretty painless. There is still the question of what something like this would do:

scripted_linear = torch.jit.script(nn.Linear(10, 10)) isinstance(scripted_linear, nn.Linear) # true or false?

Intuitively I feel like it should be False since it's constructing a new module, but I can see people doing something like this, which could lead to confusing behavior.

class A(nn.Module): def __init__(self, mod): self.module = mod if isinstance(self.module, nn.Linear): ... A(torch.jit.script(B())

zdevito · 2019-05-29T05:43:56Z

torch/jit/__init__.py

            constants_set = set(getattr(original, "__constants__", []))
            self.__dict__["_constants_set"] = {}

+            if not hasattr(original, '_parameters'):


We need to document better exactly what happens when a we create a ScriptModule from a Module: what gets copied, what is translated in some way, what is ignored.

zdevito · 2019-05-29T05:44:32Z

torch/jit/__init__.py

                    # TODO: need to handle this more generally when non-tensor attributes added to module
                    object.__setattr__(self, name, item)
-                elif isinstance(item, Parameter) or (isinstance(item, Module) and item is not self):
+                elif item is self:


I can't follow this code here. Not sure why it is deciding to ignore some things, and not others.

zdevito · 2019-05-29T05:47:18Z

torch/jit/__init__.py

+        exported = tuple(mod.__torchscript_export__)
+    methods = methods + exported
+
    if mod in _jit_internal.weak_modules:


It is possible we no longer want to cache this. Should calling script twice on the same nn.Module return the same module? What if you mutate something in the original class (e.g. do some model surgery) and recall it?

I agree, it should make a fresh copy each time to stay with the idea of it being a "pure" operation. I think that should be in a follow up diff though.

facebook-github-bot · 2019-05-30T04:08:08Z

@driazati merged this pull request in 8d3388a.

ssnl · 2019-05-30T04:19:31Z

torch/_jit_internal.py

    return fn


+def ignore(maybe_fn=None, *, drop_on_export=False):


kw only arg is not supported in python 2

Unlanding now

…omod

pytorchbot added oncall: jit Add this issue/PR to JIT oncall triage queue module: pybind Related to our Python bindings / interactions with other Python libraries labels May 20, 2019

driazati requested review from suo and zdevito May 20, 2019 17:57

driazati mentioned this pull request May 20, 2019

[wip][jit] Recursive modules #20724

Closed

pytorchbot added the module: nn Related to torch.nn label May 23, 2019

Squashed commits

e260151

driazati force-pushed the driazati/automod branch from dd9f511 to e260151 Compare May 28, 2019 22:24

davidriazati added 3 commits May 28, 2019 15:41

Cleanup tests

36d2b20

Fix catch

9418a89

Fix catch

f65ce09

zdevito approved these changes May 29, 2019

View reviewed changes

driazati mentioned this pull request May 29, 2019

Recursive Script Follow-ups #21082

Closed

davidriazati added 2 commits May 29, 2019 12:55

Address comments

244436a

Fix tests

28e19b4

facebook-github-bot closed this in 8d3388a May 30, 2019

facebook-github-bot added the merged label May 30, 2019

ssnl reviewed May 30, 2019

View reviewed changes

driazati reopened this May 30, 2019

davidriazati added 4 commits May 30, 2019 11:16

Fix python 2 bug

19df823

Merge branch 'master' of github.com:pytorch/pytorch into driazati/aut…

4844567

…omod

Flake

30152e0

Merge

4dc9ae2

facebook-github-bot closed this in 2f4824b May 31, 2019

facebook-github-bot deleted the driazati/automod branch July 13, 2020 17:54

mruberry added the Merged label Oct 28, 2020



		def _make_strong_submodule(field, module, parent):
		if field not in parent._modules:

		return fn


		def ignore(maybe_fn=None, *, drop_on_export=False):

[jit] Add support for recursive compilation on Modules #20708

[jit] Add support for recursive compilation on Modules #20708

Uh oh!

Conversation

driazati commented May 20, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

suo commented May 20, 2019

Uh oh!

driazati commented May 20, 2019

Uh oh!

zdevito commented May 21, 2019

Uh oh!

zdevito left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented May 30, 2019

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

driazati commented May 20, 2019 •

edited

Loading