Repository: Lightning-AI/pytorch-lightning

Issue:
Could I convert lightning module to onnx? Thanks!
## 🚀 Feature
pytorch Lightning works very good, but I cannot find any comments or examples to guide my convert to onnx from a pretrained lightning model, doesn't lightning module only use for researching purpose, without support of onnx cross platform?


Discussion:
Hi! thanks for your contribution!, great first issue!
good point, @PyTorchLightning/core-contributors thoughts?
From the docs https://pytorch.org/docs/stable/onnx.html it seems easy enough. We would require users to define the `self.example_array` attribute, but that's it.

Do we already have a method for saving the model that can be extended? 
a lightningModule is just a nn.Module.

Wouldn't this just work?

```python
# Input to the model
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model(x)

# Export the model
torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "super_resolution.onnx",   # where to save the model (can be a file or 
```

But we should just make this automatic...
@stevenluzheng123456 want to submit a PR for:

```python

model = LitModel(...)
model.to_onnx(x)
```

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
@@ -2,6 +2,7 @@
 import inspect
 import os
 import re
+import tempfile
 from abc import ABC, abstractmethod
 from argparse import Namespace
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -1723,6 +1724,44 @@ def _set_hparams(self, hp: Union[dict, Namespace, str]) -> None:
         else:
             self._hparams = hp
 
+    def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwargs):
+        """Saves the model in ONNX format
+
+        Args:
+            file_path: The path of the file the model should be saved to.
+            input_sample: A sample of an input tensor for tracing.
+            **kwargs: Will be passed to torch.onnx.export function.
+
+        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()
+            ...     input_sample = torch.randn((1, 64))
+            ...     model.to_onnx(tmpfile.name, input_sample, export_params=True)
+            ...     os.path.isfile(tmpfile.name)
+            True
+        """
+
+        if isinstance(input_sample, Tensor):
+            input_data = input_sample
+        elif self.example_input_array is not None:
+            input_data = self.example_input_array
+        else:
+            raise ValueError(f'input_sample and example_input_array tensors are both missing.')
+
+        if 'example_outputs' not in kwargs:
+            self.eval()
+            kwargs['example_outputs'] = self(input_data)
+
+        torch.onnx.export(self, input_data, file_path, **kwargs)
+
     @property
     def hparams(self) -> Union[AttributeDict, str]:
         if not hasattr(self, '_hparams'):


Test patch:
diff --git a/tests/base/model_template.py b/tests/base/model_template.py
--- a/tests/base/model_template.py
+++ b/tests/base/model_template.py
@@ -73,9 +73,7 @@ def __init__(
         self.test_step_end_called = False
         self.test_epoch_end_called = False
 
-        # if you specify an example input, the summary will show input/output for each layer
-        # TODO: to be fixed in #1773
-        # self.example_input_array = torch.rand(5, 28 * 28)
+        self.example_input_array = torch.rand(5, 28 * 28)
 
         # build model
         self.__build_model()
diff --git a/tests/models/test_onnx_save.py b/tests/models/test_onnx_save.py
new file mode 100644
--- /dev/null
+++ b/tests/models/test_onnx_save.py
@@ -0,0 +1,114 @@
+import os
+
+import onnxruntime
+import pytest
+import torch
+import numpy as np
+import tests.base.develop_pipelines as tpipes
+import tests.base.develop_utils as tutils
+from pytorch_lightning import Trainer
+from tests.base import EvalModelTemplate
+
+
+def test_model_saves_with_input_sample(tmpdir):
+    """Test that ONNX model saves with input sample and size is greater than 3 MB"""
+    model = EvalModelTemplate()
+    trainer = Trainer(max_epochs=1)
+    trainer.fit(model)
+
+    file_path = os.path.join(tmpdir, "model.onxx")
+    input_sample = torch.randn((1, 28 * 28))
+    model.to_onnx(file_path, input_sample)
+    assert os.path.isfile(file_path)
+    assert os.path.getsize(file_path) > 3e+06
+
+
+def test_model_saves_with_example_output(tmpdir):
+    """Test that ONNX model saves when provided with example output"""
+    model = EvalModelTemplate()
+    trainer = Trainer(max_epochs=1)
+    trainer.fit(model)
+
+    file_path = os.path.join(tmpdir, "model.onxx")
+    input_sample = torch.randn((1, 28 * 28))
+    model.eval()
+    example_outputs = model.forward(input_sample)
+    model.to_onnx(file_path, input_sample, example_outputs=example_outputs)
+    assert os.path.exists(file_path) is True
+
+
+def test_model_saves_with_example_input_array(tmpdir):
+    """Test that ONNX model saves with_example_input_array and size is greater than 3 MB"""
+    model = EvalModelTemplate()
+    file_path = os.path.join(tmpdir, "model.onxx")
+    model.to_onnx(file_path)
+    assert os.path.exists(file_path) is True
+    assert os.path.getsize(file_path) > 3e+06
+
+
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_model_saves_on_multi_gpu(tmpdir):
+    """Test that ONNX model saves on a distributed backend"""
+    tutils.set_random_master_port()
+
+    trainer_options = dict(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        limit_train_batches=10,
+        limit_val_batches=10,
+        gpus=[0, 1],
+        distributed_backend='ddp_spawn',
+        progress_bar_refresh_rate=0
+    )
+
+    model = EvalModelTemplate()
+
+    tpipes.run_model_test(trainer_options, model)
+
+    file_path = os.path.join(tmpdir, "model.onxx")
+    model.to_onnx(file_path)
+    assert os.path.exists(file_path) is True
+
+
+def test_verbose_param(tmpdir, capsys):
+    """Test that output is present when verbose parameter is set"""
+    model = EvalModelTemplate()
+    file_path = os.path.join(tmpdir, "model.onxx")
+    model.to_onnx(file_path, verbose=True)
+    captured = capsys.readouterr()
+    assert "graph(%" in captured.out
+
+
+def test_error_if_no_input(tmpdir):
+    """Test that an exception is thrown when there is no input tensor"""
+    model = EvalModelTemplate()
+    model.example_input_array = None
+    file_path = os.path.join(tmpdir, "model.onxx")
+    with pytest.raises(ValueError, match=r'input_sample and example_input_array tensors are both missing'):
+        model.to_onnx(file_path)
+
+
+def test_if_inference_output_is_valid(tmpdir):
+    """Test that the output inferred from ONNX model is same as from PyTorch"""
+    model = EvalModelTemplate()
+    trainer = Trainer(max_epochs=5)
+    trainer.fit(model)
+
+    model.eval()
+    with torch.no_grad():
+        torch_out = model(model.example_input_array)
+
+    file_path = os.path.join(tmpdir, "model.onxx")
+    model.to_onnx(file_path, model.example_input_array, export_params=True)
+
+    ort_session = onnxruntime.InferenceSession(file_path)
+
+    def to_numpy(tensor):
+        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
+
+    # compute ONNX Runtime output prediction
+    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(model.example_input_array)}
+    ort_outs = ort_session.run(None, ort_inputs)
+
+    # compare ONNX Runtime and PyTorch results
+    assert np.allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)