[jit] Implement more of of the nn.Module API #28828

zdevito · 2019-10-29T05:47:09Z

Stack from ghstack:

[jit] Implement more of of the nn.Module API #28828 [jit] Implement more of of the nn.Module API

This updates torch::script::Module to more closely match the behavior
of nn.Module. In particular, it implements the (optionally recurisive)
iterators that retrieve submodules, parameters, and buffers and makes
their names match the python versions.

This also removes the individual accessors for Parameter, Module, Buffer, etc.
and replaces them with a single attr function which is equivalent to
writing a.foo in Python (setattr emulates a.foo = v).
As we build out the user-facing API for TorchScript values this will end
up matching how an attribute is accessed on general objects.

This PR preservers the python bindings for script::Module by emulating the
old API at the binding level. A followup will clean up the usage to more
directly match the C++ API.

Differential Revision: D18197611

This updates torch::script::Module to more closely match the behavior of nn.Module. In particular, it implements the (optionally recurisive) iterators that retrieve submodules, parameters, and buffers and makes their names match the python versions. This also removes the individual accessors for Parameter, Module, Buffer, etc. and replaces them with a single `attr` function which is equivalent to writing `a.foo` in Python (`setattr` emulates `a.foo = v`). As we build out the user-facing API for TorchScript values this will end up matching how an attribute is accessed on general objects. This PR preservers the python bindings for script::Module by emulating the old API at the binding level. A followup will clean up the usage to more directly match the C++ API.

ZolotukhinM

I have two comments:

It seems that the old API for accessing attributes had checks for the type of the attribute (module/buffer/parameter), while the new API returns a generic one. We probably need to add these removed checks to the places where it was used to preserve the behavior (I commented in one of such sites, but there are probably more).
The code for recursive iterators is hard to understand. I know what it's supposed to be doing but it's difficult to follow it even with that knowledge. Can we please add some comments (and update the old comments)? Some classes like Policy et al would also benefit from brief comments.

ZolotukhinM · 2019-10-29T18:05:09Z

torch/csrc/api/src/serialize/input-archive.cpp

+  if (!module_.hasattr(key)) {
    return false;
  }
+  archive.module_ = module_.attr(key).toModule();


Wouldn't we crash here if the attribute is there, but it's not a module? In the original code we returned false in this case.

Yeah, I can make this more specific.

torch/csrc/jit/script/module.h

ZolotukhinM · 2019-10-29T18:17:10Z

torch/csrc/jit/script/module.h

-struct NameValue {
-  std::string name;
-  IValue value;
+struct Frame {


The name Frame is already overloaded in various contexts, are we sure we want to use it here as well?

Ill change it to something less ambiguous.

[jit] Implement more of of the nn.Module API This updates torch::script::Module to more closely match the behavior of nn.Module. In particular, it implements the (optionally recurisive) iterators that retrieve submodules, parameters, and buffers and makes their names match the python versions. This also removes the individual accessors for Parameter, Module, Buffer, etc. and replaces them with a single `attr` function which is equivalent to writing `a.foo` in Python (`setattr` emulates `a.foo = v`). As we build out the user-facing API for TorchScript values this will end up matching how an attribute is accessed on general objects. This PR preservers the python bindings for script::Module by emulating the old API at the binding level. A followup will clean up the usage to more directly match the C++ API. gh-metadata: pytorch pytorch 28828 gh/zdevito/130/head

ZolotukhinM

Overall looks good, please find some comments inline.

ZolotukhinM · 2019-10-30T22:23:24Z

torch/csrc/jit/passes/quantization.cpp

  script::Module observer = observer_module.clone();
  std::string observer_name = "_observer_" + std::to_string(uid_++);
-  while (module.find_module(observer_name)) {
+  while (module.hasattr(observer_name)) {


Ouch, there was a bug previously! Good thing that it's gonna be fixed now.

ZolotukhinM · 2019-10-30T22:27:16Z

torch/csrc/jit/passes/quantization.cpp

    // Queue submodules for processing
-    for (const script::NameModule& submodule : current.get_modules()) {
-      worklist.push(submodule.module);
+    for (const script::NameModule& submodule : current.named_children()) {


Nit: we can probably use children instead of named_children here.

ZolotukhinM · 2019-10-30T22:28:00Z

torch/csrc/jit/passes/quantization.cpp

    InsertPrepackUnpack(graph);
-    for (script::NameModule m : module.get_modules()) {
-      InsertPrepackUnpack(m.module);
+    for (script::NameModule m : module.named_children()) {


ZolotukhinM · 2019-10-30T22:28:34Z

torch/csrc/jit/passes/quantization.cpp

    FoldPrepackedWeightIntoModule(
        module, method.name(), linear_params_module, conv_params_module);
-    for (script::NameModule m : module.get_modules()) {
+    for (script::NameModule m : module.named_children()) {


ZolotukhinM · 2019-10-30T22:31:50Z

torch/csrc/jit/script/init.cpp

+  slot_dict_impl<Policy>(self.module_object()).setattr(name, std::move(value));
+}
+
+static py::object get_generic(Module& self, const std::string& name) {


Should this also be templatized by Policy?

My next patch is going to remove all of those and replace them with attr, so I didn't bother with it here. All of these are private methods in our ScriptModule implementation that are already guarded.

ZolotukhinM · 2019-10-30T23:03:27Z

torch/csrc/jit/script/module.h

+      IValue v) {
+    std::string name;
+    if (frames.size() == 1) {
+      name = (frames.back().i_ == -1) ? "" : nameFragment(frames.back());


Do we really have to special case size==1? It looks like it can be handled just fine by the loop below.

It's to avoid the overhead of allocating a ostringstream and copying the string twice in the very common non-recursive case.

ZolotukhinM · 2019-10-30T23:03:57Z

torch/csrc/jit/script/module.h

+        if (i > 0) {
+          ss << ".";
+        }
+        ss << nameFragment(frames[i]);


What if frames[i].i_ == -1? Should we assert that it's not the case?

The getAttributeName is going to assert in that case, so I didn't add another one. It's not a bug because only the top-level frame can have this.

ZolotukhinM · 2019-10-30T23:05:46Z

torch/csrc/jit/script/module.h

-           (type_ && module_.entity_type(i_) != *type_)) {
-      ++i_;
+  // return_module() is a corner case where instead of returning a submodule
+  // of root, we are return root itself, because we are iterating modules(),


Typo: "are return"

ZolotukhinM · 2019-10-30T23:06:38Z

torch/csrc/jit/script/module.h

+  // return_module() is a corner case where instead of returning a submodule
+  // of root, we are return root itself, because we are iterating modules(),
+  // which contains the root module itself.
+  // It is represented with a single Frame whose index is -1.


Nit: since Frame was renamed to SlotCursor, we need to update all its references accordingly.

ZolotukhinM · 2019-10-30T23:09:43Z

torch/csrc/jit/script/module.h

+    }
+    // the last traversal action advanced beyond the number of slots in the
+    // module so continue the iteration in the parent.
+    if (top().i_ >= int64_t(top().module_.num_slots())) {


Should we do it while we're beyond the number of slots instead of if? I.e. shall we pop back until we reach a valid position or do we intentionally want to pop back only once?

That would be correct too, but I went with this way to make the components easier to understand: next() does 1 step, while_not_valid is responsible for repeating.

[jit] Implement more of of the nn.Module API This updates torch::script::Module to more closely match the behavior of nn.Module. In particular, it implements the (optionally recurisive) iterators that retrieve submodules, parameters, and buffers and makes their names match the python versions. This also removes the individual accessors for Parameter, Module, Buffer, etc. and replaces them with a single `attr` function which is equivalent to writing `a.foo` in Python (`setattr` emulates `a.foo = v`). As we build out the user-facing API for TorchScript values this will end up matching how an attribute is accessed on general objects. This PR preservers the python bindings for script::Module by emulating the old API at the binding level. A followup will clean up the usage to more directly match the C++ API. gh-metadata: pytorch pytorch 28828 gh/zdevito/130/head

kostmo · 2019-10-31T22:42:53Z

CircleCI build failures summary

As of commit 74bc5c4:

0/2 recognized as flaky
2/2 broken upstream. You may want to rebase on the latest viable branch.
0/2 failures introduced in this PR

Here are the reasons each build failed.

This comment was automatically generated by Dr. CI.
Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker.

This comment has been revised 17 time(s).

[jit] Implement more of of the nn.Module API This updates torch::script::Module to more closely match the behavior of nn.Module. In particular, it implements the (optionally recurisive) iterators that retrieve submodules, parameters, and buffers and makes their names match the python versions. This also removes the individual accessors for Parameter, Module, Buffer, etc. and replaces them with a single `attr` function which is equivalent to writing `a.foo` in Python (`setattr` emulates `a.foo = v`). As we build out the user-facing API for TorchScript values this will end up matching how an attribute is accessed on general objects. This PR preservers the python bindings for script::Module by emulating the old API at the binding level. A followup will clean up the usage to more directly match the C++ API. gh-metadata: pytorch pytorch 28828 gh/zdevito/130/head

Summary: Pull Request resolved: pytorch/pytorch#28828 This updates torch::script::Module to more closely match the behavior of nn.Module. In particular, it implements the (optionally recurisive) iterators that retrieve submodules, parameters, and buffers and makes their names match the python versions. This also removes the individual accessors for Parameter, Module, Buffer, etc. and replaces them with a single `attr` function which is equivalent to writing `a.foo` in Python (`setattr` emulates `a.foo = v`). As we build out the user-facing API for TorchScript values this will end up matching how an attribute is accessed on general objects. This PR preservers the python bindings for script::Module by emulating the old API at the binding level. A followup will clean up the usage to more directly match the C++ API. Test Plan: Imported from OSS Differential Revision: D18197611 Pulled By: zdevito fbshipit-source-id: 7ee4dcbb258605d1c988314b05d938423f1ccee5

facebook-github-bot · 2019-11-07T11:05:54Z

@zdevito merged this pull request in 7963631.

pietern · 2019-11-07T13:08:43Z

@zdevito This raced with #29208 and broke master.

Working on a fix instead of reverting because I think it's an easy fix.

pietern · 2019-11-07T13:24:51Z

Not an easy fix after all.

Reverting #29208 now.

Test Plan: revert-hammer Differential Revision: D18350353 Original commit changeset: 2026c8ab7650 fbshipit-source-id: 401f34cb276c3ea34a5439de4c3415969a04ab2a

zdevito requested review from apaszke, ebetica, goldsborough and yf225 as code owners October 29, 2019 05:47

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Oct 29, 2019

zdevito requested a review from ZolotukhinM October 29, 2019 05:48

ZolotukhinM reviewed Oct 29, 2019

View reviewed changes

zdevito added 3 commits October 29, 2019 16:31

zdevito requested a review from ZolotukhinM October 30, 2019 21:53

ZolotukhinM approved these changes Oct 30, 2019

View reviewed changes

zdevito added 2 commits October 30, 2019 17:25

zdevito added 6 commits November 4, 2019 10:37

jerryzh168 mentioned this pull request Nov 6, 2019

[fix] clone should preserve the type of attribute #29269

Closed

facebook-github-bot closed this in 7963631 Nov 7, 2019

facebook-github-bot added the merged label Nov 7, 2019

pietern mentioned this pull request Nov 7, 2019

dump operator names of a module and its sub-modules. #29208

Closed

pietern referenced this pull request Nov 7, 2019

Revert D18350353: dump operator names of a module and its sub-modules.

78a34d3

Test Plan: revert-hammer Differential Revision: D18350353 Original commit changeset: 2026c8ab7650 fbshipit-source-id: 401f34cb276c3ea34a5439de4c3415969a04ab2a

facebook-github-bot deleted the gh/zdevito/130/head branch November 10, 2019 15:16

eellison mentioned this pull request Nov 20, 2019

Move the attributes of a module to the given device #29987

Open

wanchaol mentioned this pull request Oct 19, 2020

[jit] Support named_parameters in TorchScript #46555

Closed

mruberry added the Merged label Oct 28, 2020

[jit] Implement more of of the nn.Module API #28828

[jit] Implement more of of the nn.Module API #28828

Uh oh!

Conversation

zdevito commented Oct 29, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ZolotukhinM left a comment

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!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZolotukhinM left a comment

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!

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!

kostmo commented Oct 31, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CircleCI build failures summary

Uh oh!

facebook-github-bot commented Nov 7, 2019

Uh oh!

pietern commented Nov 7, 2019

Uh oh!

pietern commented Nov 7, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

zdevito commented Oct 29, 2019 •

edited

Loading

kostmo commented Oct 31, 2019 •

edited

Loading