to_torchscript method for LightningModule

[awaelchli]

What does this PR do?

Fixes #3080

Tests:

• tests exporting several modules (GAN, RNN, Standard test model)
• test save and load
• test that unused properties like datamodules etc, and hooks get removed from scriptmodule

NOTE:

• Renamed TestGAN because in pytest, classes named Test-... are treated as pytest-TestClasses which are supposed to hold testing methods and this messes with torchscript somehow.
• Had to skip some tests because torch.jit.save crashes on pytorch <= 1.4. I assume it's a bug that got fixed, I don't understand the cause though.

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Did you verify new and existing tests pass locally with your changes?
• If you made a notable change (that affects users), did you update the CHANGELOG?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #3258 into master will decrease coverage by 6%.
The diff coverage is 90%.

@@           Coverage Diff           @@
##           master   #3258    +/-   ##
=======================================
- Coverage      86%     80%    -6%     
=======================================
  Files          88      90     +2     
  Lines        8132    8995   +863     
=======================================
+ Hits         6962    7185   +223     
- Misses       1170    1810   +640     

[ananthsub]

Super exciting! Thanks for working on this @awaelchli

[pep8speaks]

Hello @awaelchli! Thanks for updating this PR.

• In the file tests/base/models.py:

Line 70:121: E501 line too long (122 > 120 characters)

Comment last updated at 2020-09-03 16:36:13 UTC

[justusschock]

LGTM.

I only have one question: Theoretically it is possible to also train ScriptModules. But this PR does not enable this, right? Because that way saving and loading may have to properly deal with the attributes.

Not sure, what happens if you derive your model from LightningModule and ScriptModule though

[awaelchli]

@justusschock

I only have one question: Theoretically it is possible to also train ScriptModules. But this PR does not enable this, right? Because that way saving and loading may have to properly deal with the attributes.

my understanding is scripting removes all properties, methods etc. that do not appear in forward (I assert that in thet test actually). all the lightning stuff is gone. haven't tried, but if you were to train a script module, it would have to use forward, so pure pytorch training loop.

we talked with @ananthsub of a way to load a script module into a LightningModule, that would essentially replace the forward and then could be trained. if we can formulate a concrete plan for how to integrate that, I'm happy to follow up on it.

[rohitgr7]

NICE!!