[jit] Make `nn.Transformer` TorchScript compatible #28561

driazati · 2019-10-23T23:34:11Z

This makes nn.Transformer usable from TorchScript. It preserves backwards compatibility via __setstate__ on the encoder/decoder.

Fixes #24173

Differential Revision: D18124753

This makes MultiheadedAttention TorchScript compatible. It changes the BC-compatible code in `forward` to use `__setstate__` instead so that `torch.load` still works correctly for old models

…nto driazati/transformer/2

zhangguanheng66 · 2019-10-25T17:24:47Z

Add BC breaking tag for release note. @gchanan

gchanan · 2019-10-27T00:24:58Z

Remember to add label "topic: bc-breaking" to such changes as well :).

zhangguanheng66 · 2019-10-28T17:07:46Z

torch/nn/modules/transformer.py

@@ -171,11 +173,10 @@ def forward(self, src, mask=None, src_key_padding_mask=None):
        """
        output = src

-        for i in range(self.num_layers):


So range is not jitable as well?

ModuleLists can't be indexed, but can be used with for-in loops. This should be the same

zhangguanheng66 · 2019-10-28T17:08:43Z

torch/nn/modules/transformer.py


-        if self.norm:
+        if self.norm is not None:


Is this change due to jitable?

Right the check has to be that the value is not None so this branch doesn't get compiled if there is no code to call (i.e. if norm is None)

zhangguanheng66 · 2019-10-28T17:10:22Z

torch/nn/modules/transformer.py


    def __init__(self, decoder_layer, num_layers, norm=None):
        super(TransformerDecoder, self).__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
-        self.num_layers = num_layers


This is a BC breaking

This seems like redundant info since someone can just do len(self.layers) to get the same value. Since this change is BC-breaking anyways I think this cleanup is warranted

Since I want to cover BC for transformer, please add those two attributes back and I will approve it.

self.num_layers is added to Encoder but not Decoder.

zhangguanheng66

Do you need tests to cover the scriptable transformer?

…nsformer/2

added test

zhangguanheng66

We are thinking about the _load_from_state_dict func to avoid BC breaking. See the PR I added. #29001

driazati · 2019-11-04T21:14:16Z

By the time 1.4 is released there will have been a whole minor version for people to update their models, which I think is enough time to make a BC-breaking change here on loading old models.

Either way it'd be good to settle on something, the first version of these PRs (this and #28555) had the proper BC-maintaining changes which I reverted based on our discussions.

zhangguanheng66 · 2019-11-04T21:52:45Z

By the time 1.4 is released there will have been a whole minor version for people to update their models, which I think is enough time to make a BC-breaking change here on loading old models.

Either way it'd be good to settle on something, the first version of these PRs (this and #28555) had the proper BC-maintaining changes which I reverted based on our discussions.

Yes. Based on our discussions, we agreed to remove hasattr func, which introduces BC breaking. Adding _load_from_state_dict could avoid such breaking while your current version still works.

Since I have the PR (#28555) to add _load_from_state_dict func for MultiheadAttention, I will add _load_from_state_dict func to nn.Transformer there, as well.

zhangguanheng66 · 2019-11-20T14:38:54Z

Please consider to overwrite __setstate__ func to avoid BC breaking. #29001.
Once it's done, we should remove the BC-breaking flag.

driazati · 2019-11-27T01:02:59Z

torch/nn/modules/transformer.py

@@ -355,10 +366,7 @@ def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
-        if hasattr(self, "activation"):


Since this gets patched up in __setstate__, the other case is not necessary

zhangguanheng66 · 2019-11-27T02:06:46Z

torch/nn/modules/transformer.py


    def __init__(self, encoder_layer, num_layers, norm=None):
        super(TransformerEncoder, self).__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
-        self.num_layers = num_layers


nit: this will be a BC breaking. Prefer not to remove it since we won't claim a BC breaking for this module. Will this block a jitable transformer?

I re-added it in the latest commit

zhangguanheng66

One more comment. Thanks.

torch/nn/modules/transformer.py

zhangguanheng66

self.num_layers is added to Encoder but not Decoder. Please add it to Decoder and the PR is ready to merge. Thanks.

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 · 2019-12-11T19:53:10Z

@driazati merged this pull request in 1f87e82.

Landed in #28561

Summary: This makes `nn.Transformer` usable from TorchScript. It preserves backwards compatibility via `__setstate__` on the encoder/decoder. Fixes pytorch#24173 Pull Request resolved: pytorch#28561 Differential Revision: D18124753 Pulled By: driazati fbshipit-source-id: 7314843e5aa9c9bf974c4672e4edb24ed8ef4a6f

Your Name added 2 commits October 23, 2019 15:00

[jit] Support MultiheadedAttention module

1ea1bae

This makes MultiheadedAttention TorchScript compatible. It changes the BC-compatible code in `forward` to use `__setstate__` instead so that `torch.load` still works correctly for old models

[jit] Make nn.Transformer TorchScript compatible

2c16bf8

driazati requested a review from apaszke as a code owner October 23, 2019 23:34

Your Name added 11 commits October 23, 2019 16:35

update

64eb281

update

74daf9c

fix flake

446e929

fix, add test

69c4246

Add in_proj_weight as none param

3e1d9f7

Add in_proj_weight as none param

92759e8

Remove set state

370dfcb

Merge branch 'driazati/transformer/1' of github.com:pytorch/pytorch i…

c85fb3a

…nto driazati/transformer/2

Remove activation checks

d75fe60

Remove activation checks

2d3b222

Remove activation checks

939863b

driazati mentioned this pull request Oct 24, 2019

[jit] Support MultiheadedAttention module #28555

Closed

zhangguanheng66 changed the title ~~[jit] Make nn.Transformer TorchScript compatible~~ [jit][BC-breaking] Make nn.Transformer TorchScript compatible Oct 25, 2019

gchanan added the module: bc-breaking Related to a BC-breaking change label Oct 27, 2019

driazati requested a review from zhangguanheng66 October 28, 2019 15:58

zhangguanheng66 reviewed Oct 28, 2019

View reviewed changes

zhangguanheng66 previously requested changes Oct 28, 2019

View reviewed changes

driazati requested a review from zhangguanheng66 October 30, 2019 00:54

driazati changed the base branch from driazati/transformer/1 to master October 30, 2019 00:55

driazati changed the base branch from master to driazati/transformer/1 October 30, 2019 00:55

zhangguanheng66 reviewed Oct 31, 2019

View reviewed changes

Your Name added 2 commits November 4, 2019 10:16

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

bcfa1b6

…nsformer/2

Add test

bdae4f5

driazati requested a review from ebetica as a code owner November 4, 2019 19:50

driazati requested a review from zhangguanheng66 November 4, 2019 19:50

zhangguanheng66 reviewed Nov 4, 2019

View reviewed changes

Your Name added 2 commits November 26, 2019 16:53

Merge

2daa929

Re-add BC compat fixes

ba358ab

driazati commented Nov 27, 2019

View reviewed changes

driazati changed the title ~~[jit][BC-breaking] Make nn.Transformer TorchScript compatible~~ [jit] Make nn.Transformer TorchScript compatible Nov 27, 2019

driazati requested a review from zhangguanheng66 November 27, 2019 01:03

zhangguanheng66 reviewed Nov 27, 2019

View reviewed changes

re-add num_layers

cf28321

zhangguanheng66 requested changes Nov 27, 2019

View reviewed changes

torch/nn/modules/transformer.py Outdated Show resolved Hide resolved

Fix activation default

f844b3f

driazati removed the module: bc-breaking Related to a BC-breaking change label Nov 28, 2019

driazati requested a review from zhangguanheng66 November 28, 2019 00:21

zhangguanheng66 approved these changes Dec 9, 2019

View reviewed changes

Your Name added 2 commits December 9, 2019 14:13

Fix missing num_layers

a057137

Fix mypy

c75eb5e

facebook-github-bot reviewed Dec 10, 2019

View reviewed changes

facebook-github-bot closed this in 1f87e82 Dec 11, 2019

facebook-github-bot added the merged label Dec 11, 2019

driazati pushed a commit that referenced this pull request Dec 17, 2019

[jit] Make nn.Transformer TorchScript compatible

f8a24e9

Landed in #28561

driazati mentioned this pull request Dec 17, 2019

[jit] Make nn.Transformer TorchScript compatible #31390

Closed

facebook-github-bot deleted the driazati/transformer/2 branch July 13, 2020 17:55

mruberry added the Merged label Oct 28, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[jit] Make `nn.Transformer` TorchScript compatible #28561

[jit] Make `nn.Transformer` TorchScript compatible #28561

driazati commented Oct 23, 2019 •

edited

Loading

zhangguanheng66 commented Oct 25, 2019

gchanan commented Oct 27, 2019

zhangguanheng66 Oct 28, 2019

driazati Oct 30, 2019

zhangguanheng66 Oct 28, 2019

driazati Oct 30, 2019

zhangguanheng66 Oct 28, 2019

driazati Oct 30, 2019

zhangguanheng66 Nov 4, 2019 •

edited

Loading

zhangguanheng66 Dec 9, 2019

zhangguanheng66 left a comment

zhangguanheng66 left a comment

driazati commented Nov 4, 2019 •

edited

Loading

zhangguanheng66 commented Nov 4, 2019

zhangguanheng66 commented Nov 20, 2019

driazati Nov 27, 2019

zhangguanheng66 Nov 27, 2019

driazati Nov 27, 2019

zhangguanheng66 left a comment

zhangguanheng66 left a comment

facebook-github-bot left a comment

facebook-github-bot commented Dec 11, 2019

[jit] Make nn.Transformer TorchScript compatible #28561

[jit] Make nn.Transformer TorchScript compatible #28561

Conversation

driazati commented Oct 23, 2019 • edited Loading

zhangguanheng66 commented Oct 25, 2019

gchanan commented Oct 27, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhangguanheng66 Nov 4, 2019 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhangguanheng66 left a comment

Choose a reason for hiding this comment

zhangguanheng66 left a comment

Choose a reason for hiding this comment

driazati commented Nov 4, 2019 • edited Loading

zhangguanheng66 commented Nov 4, 2019

zhangguanheng66 commented Nov 20, 2019

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhangguanheng66 left a comment

Choose a reason for hiding this comment

zhangguanheng66 left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot commented Dec 11, 2019

[jit] Make `nn.Transformer` TorchScript compatible #28561

[jit] Make `nn.Transformer` TorchScript compatible #28561

driazati commented Oct 23, 2019 •

edited

Loading

zhangguanheng66 Nov 4, 2019 •

edited

Loading

driazati commented Nov 4, 2019 •

edited

Loading