Repository: Lightning-AI/pytorch-lightning

Issue:
Support a to_torchscript function on the LightningModule
## 🚀 Feature
Support a conversion function to [PyTorch JIT](https://pytorch.org/docs/stable/jit.html) similar to what's available for [ONNX](https://pytorch-lightning.readthedocs.io/en/latest/lightning-module.html#to-onnx).
 
### Motivation
TorchScript is a way to create serializable and optimizable models from PyTorch code. Any TorchScript program can be saved from a Python process and loaded in a process where there is no Python dependency. TorchScript is a method by which users can serve PyTorch models efficiently. 

### Pitch
By default, we can use TorchScript to script the LightningModule. Users can override this in their own lightning modules to use tracing, or to script specific nn.Modules inside their LightningModule. This can then be extended to other Lightning utilities like model checkpointing, so we can save TorchScript or ONNX converted models alongside the best model checkpoints to make going to serving even easier to do

```
    def to_torchscript(self):
        """Saves the model as a JIT module. 
            This can be overridden to support custom TorchScript module export
        Example:
            >>> class SimpleModel(LightningModule):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self.l1 = torch.nn.Linear(in_features=64, out_features=4)
            ...
            ...     def forward(self, x):
            ...         return torch.relu(self.l1(x.view(x.size(0), -1)))
            >>> with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
            ...     model = SimpleModel()
            ...     torch.jit.save(model.to_torchscript(), tmpfile.name)
            ...     os.path.isfile(tmpfile.name)
            True
        """
       return torch.jit.script(self.eval())
```





Discussion:
I like this proposal and it shall be quite straight forward, mind send a PR? 🐰 
cc: @PyTorchLightning/core-contributors 
@Borda regarding the API, what if someone wanted to expose multiple torchscript modules from their lightning module. E.g. their lightning module does GAN training, and they want to export both the generator and discriminator separately. Should the interface allow for multiple scripted modules to be returned?
I think that would be super useful in these cases! How about something like
```python
model.to_torchscript()  # exports whole module
model.to_torchscript(children=["discriminator", "generator"])  # returns the list of script modules
```
or should that just be part of the override the user can do, without any argument passed to the method?
> I think that would be super useful in these cases! How about something like
> 
> ```python
> model.to_torchscript()  # exports whole module
> model.to_torchscript(children=["discriminator", "generator"])  # returns the list of script modules
> ```
> 
> or should that just be part of the override the user can do, without any argument passed to the method?

so, the PL module would need attrs with discriminator and generator no? i think this makes sense as long as we specify that.

@ananthsub what about transforms? or a data pipeline the user might also need when exporting to torchscript/onnx. how are you thinking about that? 
Yeah, the children paraneter is good =) 
> @ananthsub what about transforms? or a data pipeline the user might also need when exporting to torchscript/onnx. how are you thinking about that?

Yes, combining with transforms is an important need. In order to serve models, we will also extract the transforms from the data module and "glue" them into the forward function (assuming the transforms are also TorchScriptable)

> so, the PL module would need attrs with discriminator and generator no? i think this makes sense as long as we specify that.

I think it'd be easier to define these on the pl_module outside of this function. I'm hesitant to make this take arbitrary args and add magic inside of the base call here. What do you think about this proposal? 

I think an API like this would be extensible. We are incubating this via a Lightning callback which people can extend

```
    def to_torchscript(self) -> Dict[str, nn.Module]:
        # By default, scripts the whole lightning module and wraps this in a dictionary from name to scripted model
        # if you would like to customize the modules that are scripted
        # or you want to use tracing you should override this method
        with pl_module.eval() and torch.no_grad():
            scripted_model = torch.jit.script(pl_module)
        return {DEFAULT_MODEL_FILENAME: scripted_model}
```

For consistency's sake, pointing out that to_onnx() takes a file path, though I would much rather prefer this proposed solution of returning Dict[str, nn.Module].
> Yes, combining with transforms is an important need. In order to serve models, we will also extract the transforms from the data module and "glue" them into the forward function (assuming the transforms are also TorchScriptable)

there are train_transforms, val_transforms and test_transforms in data module. Could we assume test_transforms always be the one to torchscript with model? Sometimes, the transform used in training/testing might be different than in inference, we might need additional logic to patch the transform during export.

How about add "to_torchscript()" to data module? such that model doesn't need to know the members in data module and any special patching for transform happens inside of data module.
@hudeven I think that datamodule shall not be part of the exported model as we move to data-agnostic models...
thoughts @nateraw @williamFalcon 
> @hudeven I think that datamodule shall not be part of the exported model as we move to data-agnostic models...
thoughts @nateraw @williamFalcon

That makes sense. I think we can go with the simplest approach relying on just the 

It'd be great if we could also bake in bundled inputs into the torchscript export, as this would be broadly useful: https://github.com/pytorch/pytorch/blob/master/torch/utils/bundled_inputs.py#L17
@ananthsub we also have the option to use torch.jit.trace. Should we support both in the same function, or better not?
I was thinking `to_torchscript(trace=False)`?

Resolving patch:
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -23,7 +23,7 @@
 
 import torch
 import torch.distributed as torch_distrib
-from torch import Tensor
+from torch import Tensor, ScriptModule
 from torch.nn import Module
 from torch.nn.parallel import DistributedDataParallel
 from torch.optim.optimizer import Optimizer
@@ -184,6 +184,7 @@ def forward(self, batch):
                     return logits
 
         """
+        return super().forward(*args, **kwargs)
 
     def training_step(self, *args, **kwargs):
         r"""
@@ -1729,6 +1730,54 @@ def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwarg
 
         torch.onnx.export(self, input_data, file_path, **kwargs)
 
+    def to_torchscript(self, file_path: Optional[str] = None, **kwargs) -> Union[ScriptModule, Dict[str, ScriptModule]]:
+        """
+        By default compiles the whole model to a :class:`~torch.jit.ScriptModule`.
+        If you would like to customize the modules that are scripted or you want to use tracing
+        you should override this method. In case you want to return multiple modules, we
+        recommend using a dictionary.
+
+        Args:
+            file_path: Path where to save the torchscript. Default: None (no file saved).
+            **kwargs: Additional arguments that will be passed to the :func:`torch.jit.save` function.
+
+        Note:
+            - Requires the implementation of the
+              :meth:`~pytorch_lightning.core.lightning.LightningModule.forward` method.
+            - The exported script will be set to evaluation mode.
+            - It is recommended that you install the latest supported version of PyTorch
+              to use this feature without limitations. See also the :mod:`torch.jit`
+              documentation for supported features.
+
+        Example:
+            >>> class SimpleModel(LightningModule):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...         self.l1 = torch.nn.Linear(in_features=64, out_features=4)
+            ...
+            ...     def forward(self, x):
+            ...         return torch.relu(self.l1(x.view(x.size(0), -1)))
+            ...
+            >>> model = SimpleModel()
+            >>> torch.jit.save(model.to_torchscript(), "model.pt")  # doctest: +SKIP
+            >>> os.path.isfile("model.pt")  # doctest: +SKIP
+            True
+
+        Return:
+            This LightningModule as a torchscript, regardless of whether file_path is
+            defined or not.
+        """
+
+        mode = self.training
+        with torch.no_grad():
+            scripted_module = torch.jit.script(self.eval(), **kwargs)
+        self.train(mode)
+
+        if file_path is not None:
+            torch.jit.save(scripted_module, file_path)
+
+        return scripted_module
+
     @property
     def hparams(self) -> Union[AttributeDict, str]:
         if not hasattr(self, '_hparams'):


Test patch:
diff --git a/tests/base/models.py b/tests/base/models.py
--- a/tests/base/models.py
+++ b/tests/base/models.py
@@ -18,7 +18,7 @@
 
 
 class Generator(nn.Module):
-    def __init__(self, latent_dim: tuple, img_shape: tuple):
+    def __init__(self, latent_dim: int, img_shape: tuple):
         super().__init__()
         self.img_shape = img_shape
 
@@ -64,10 +64,10 @@ def forward(self, img):
         return validity
 
 
-class TestGAN(LightningModule):
+class BasicGAN(LightningModule):
     """Implements a basic GAN for the purpose of illustrating multiple optimizers."""
 
-    def __init__(self, hidden_dim, learning_rate, b1, b2, **kwargs):
+    def __init__(self, hidden_dim: int = 128, learning_rate: float = 0.001, b1: float = 0.5, b2: float = 0.999, **kwargs):
         super().__init__()
         self.hidden_dim = hidden_dim
         self.learning_rate = learning_rate
@@ -163,6 +163,7 @@ def __init__(self):
         super().__init__()
         self.rnn = nn.LSTM(10, 20, batch_first=True)
         self.linear_out = nn.Linear(in_features=20, out_features=5)
+        self.example_input_array = torch.rand(2, 3, 10)
 
     def forward(self, x):
         seq, last = self.rnn(x)
@@ -189,6 +190,7 @@ def __init__(self):
         self.c_d1_bn = nn.BatchNorm1d(128)
         self.c_d1_drop = nn.Dropout(0.3)
         self.c_d2 = nn.Linear(in_features=128, out_features=10)
+        self.example_input_array = torch.rand(2, 1, 28, 28)
 
     def forward(self, x):
         x = x.view(x.size(0), -1)
diff --git a/tests/models/test_horovod.py b/tests/models/test_horovod.py
--- a/tests/models/test_horovod.py
+++ b/tests/models/test_horovod.py
@@ -13,7 +13,7 @@
 import tests.base.develop_utils as tutils
 from pytorch_lightning import Trainer
 from tests.base import EvalModelTemplate
-from tests.base.models import TestGAN
+from tests.base.models import BasicGAN
 
 try:
     from horovod.common.util import nccl_built
@@ -145,7 +145,7 @@ def validation_step(self, batch, *args, **kwargs):
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="Horovod is not supported on Windows")
 def test_horovod_multi_optimizer(tmpdir):
-    model = TestGAN(**EvalModelTemplate.get_default_hparams())
+    model = BasicGAN(**EvalModelTemplate.get_default_hparams())
 
     # fit model
     trainer = Trainer(
diff --git a/tests/models/test_torchscript.py b/tests/models/test_torchscript.py
new file mode 100644
--- /dev/null
+++ b/tests/models/test_torchscript.py
@@ -0,0 +1,88 @@
+from distutils.version import LooseVersion
+
+import pytest
+import torch
+
+from tests.base import EvalModelTemplate
+from tests.base.datamodules import TrialMNISTDataModule
+from tests.base.models import ParityModuleRNN, BasicGAN
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    BasicGAN,
+])
+def test_torchscript_input_output(modelclass):
+    """ Test that scripted LightningModule forward works. """
+    model = modelclass()
+    script = model.to_torchscript()
+    assert isinstance(script, torch.jit.ScriptModule)
+    model.eval()
+    model_output = model(model.example_input_array)
+    script_output = script(model.example_input_array)
+    assert torch.allclose(script_output, model_output)
+
+
+@pytest.mark.parametrize("device", [
+    torch.device("cpu"),
+    torch.device("cuda", 0)
+])
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires GPU machine")
+def test_torchscript_device(device):
+    """ Test that scripted module is on the correct device. """
+    model = EvalModelTemplate().to(device)
+    script = model.to_torchscript()
+    assert next(script.parameters()).device == device
+    script_output = script(model.example_input_array.to(device))
+    assert script_output.device == device
+
+
+def test_torchscript_retain_training_state():
+    """ Test that torchscript export does not alter the training mode of original model. """
+    model = EvalModelTemplate()
+    model.train(True)
+    script = model.to_torchscript()
+    assert model.training
+    assert not script.training
+    model.train(False)
+    _ = model.to_torchscript()
+    assert not model.training
+    assert not script.training
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    BasicGAN,
+])
+def test_torchscript_properties(modelclass):
+    """ Test that scripted LightningModule has unnecessary methods removed. """
+    model = modelclass()
+    model.datamodule = TrialMNISTDataModule()
+    script = model.to_torchscript()
+    assert not hasattr(script, "datamodule")
+    assert not hasattr(model, "batch_size") or hasattr(script, "batch_size")
+    assert not hasattr(model, "learning_rate") or hasattr(script, "learning_rate")
+
+    if LooseVersion(torch.__version__) >= LooseVersion("1.4.0"):
+        # only on torch >= 1.4 do these unused methods get removed
+        assert not callable(getattr(script, "training_step", None))
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    BasicGAN,
+])
+@pytest.mark.skipif(
+    LooseVersion(torch.__version__) < LooseVersion("1.5.0"),
+    reason="torch.save/load has bug loading script modules on torch <= 1.4",
+)
+def test_torchscript_save_load(tmpdir, modelclass):
+    """ Test that scripted LightningModules is correctly saved and can be loaded. """
+    model = modelclass()
+    output_file = str(tmpdir / "model.pt")
+    script = model.to_torchscript(file_path=output_file)
+    loaded_script = torch.jit.load(output_file)
+    assert torch.allclose(next(script.parameters()), next(loaded_script.parameters()))