to_torchscript method for LightningModule

docs/source/production_inference.rst

@@ -1,3 +1,17 @@
+.. testsetup:: *
+
+    import torch
+    pytorch_lightning.core.lightning import LightningModule
+
+    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)))
+
 .. _production-inference:
 
 Inference in Production
@@ -28,3 +42,19 @@ Once you have the exported model, you can run it on your ONNX runtime in the fol
     input_name = ort_session.get_inputs()[0].name
     ort_inputs = {input_name: np.random.randn(1, 64).astype(np.float32)}
     ort_outs = ort_session.run(None, ort_inputs)
+
+
+Exporting to TorchScript
+------------------------
+
+TorchScript allows you to serialize your models in a way that it can be loaded in non-Python environments.
+The LightningModule has a handy method :meth:`~pytorch_lightning.core.lightning.LightningModule.to_torchscript`
+that returns a scripted module which you can save or directly use.
+
+.. testcode::
+
+    model = SimpleModel()
+    script = model.to_torchscript()
+
+    # save for use in production environment
+    torch.jit.save(script, "model.pt")

pytorch_lightning/core/lightning.py

@@ -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,34 @@ 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) -> torch.jit.ScriptModule:
+        """
+        Compiles the model to a :class:`~torch.jit.ScriptModule`.
+        This can be overridden to support custom TorchScript module export.
+
+        Note:
+            Requires the implementation of the :meth:`LightningModule.forward` method.
+
+        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")
+            >>> os.path.isfile("model.pt")
+            True
+        """
+        mode = self.training
+        with torch.no_grad():
+            scripted_module = torch.jit.script(self.eval())
+        self.training = mode
+        return scripted_module
+
     @property
     def hparams(self) -> Union[AttributeDict, str]:
         if not hasattr(self, '_hparams'):

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
 
@@ -67,7 +67,7 @@ def forward(self, img):
 class TestGAN(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=128, learning_rate=0.001, b1=0.5, b2=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)

tests/models/test_torchscript.py

@@ -0,0 +1,52 @@
+import pytest
+import torch
+
+from tests.base import EvalModelTemplate
+from tests.base.datamodules import TrialMNISTDataModule
+from tests.base.models import ParityModuleRNN, TestGAN
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    TestGAN,
+])
+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_output = model(model.example_input_array)
+    script_output = script(model.example_input_array)
+    assert torch.allclose(script_output, model_output)
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    TestGAN,
+])
+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")
+    assert not callable(getattr(script, "training_step", None))
+
+
+@pytest.mark.parametrize("modelclass", [
+    EvalModelTemplate,
+    ParityModuleRNN,
+    TestGAN,
+])
+def test_torchscript_save_load(tmpdir, modelclass):
+    """ Test that scripted LightningModules can be saved and loaded. """
+    model = modelclass()
+    script = model.to_torchscript()
+    assert isinstance(script, torch.jit.ScriptModule)
+    output_file = str(tmpdir / "model.pt")
+    torch.jit.save(script, output_file)
+    torch.jit.load(output_file)