Repository: Lightning-AI/pytorch-lightning
Pull request: fix ONNX model save on GPU
<!--
Please note that we have freeze state on adding new features till v1.0 release.
Anyway, any new feature suggestion is welcome and you are free to draft a PR,
 but be aware that several refactoring will take place for this major release yielding in significant collision solving.
A recommendation for feature draft is draw just outline, do not implement all details and propagate the changes to codebase...
-->

## What does this PR do?

Fixes #3144  

# Before submitting

- [x] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), Pull Request section?
- [x] 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?
- [x] Did you write any new necessary tests? 
- [x] 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](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## 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 🙃


Problem statement:
ONNX model does not save on GPU
## 🐛 Bug

Attempting to export on ONNX after training model on GPU, throws an error is the `input_sample` or `example_input_array` is not a CUDA tensor.

### To Reproduce

Steps to reproduce the behavior:

1. Train a model on GPU
2. Try to export to ONNX when  `self.example_input_array = torch.zeros(1, 1, 500, 500)` or `input_sample = torch.zeros(1, 1, 500, 500)`

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-32-cd8009a0b6a3> in <module>
      1 filepath = 'model.onnx'
----> 2 model.to_onnx(filepath, export_params=True)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/core/lightning.py in to_onnx(self, file_path, input_sample, **kwargs)
   1721         if 'example_outputs' not in kwargs:
   1722             self.eval()
-> 1723             kwargs['example_outputs'] = self(input_data)
   1724 
   1725         torch.onnx.export(self, input_data, file_path, **kwargs)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

<ipython-input-24-51cae3b5e57f> in forward(self, inputs)
     20 
     21     def forward(self, inputs):
---> 22         return self.model(inputs)
     23 
     24     def training_step(self, batch, batch_idx):

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/container.py in forward(self, input)
     98     def forward(self, input):
     99         for module in self:
--> 100             input = module(input)
    101         return input
    102 

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
    548             result = self._slow_forward(*input, **kwargs)
    549         else:
--> 550             result = self.forward(*input, **kwargs)
    551         for hook in self._forward_hooks.values():
    552             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in forward(self, input)
    351 
    352     def forward(self, input):
--> 353         return self._conv_forward(input, self.weight)
    354 
    355 class Conv3d(_ConvNd):

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    348                             _pair(0), self.dilation, self.groups)
    349         return F.conv2d(input, weight, self.bias, self.stride,
--> 350                         self.padding, self.dilation, self.groups)
    351 
    352     def forward(self, input):

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same
```


#### Code sample

filepath = 'model.onnx'
model.to_onnx(filepath, export_params=True)

### Expected behavior

Should automatically convert `example_input_array` or `input_sample` to the device type and save the model to ONNX.

Hints from issue discussion:
I would say that the problem could be the distributed way, mind check running only on a single GPU?
I ran this on Kaggle notebook. When I tried to save after training, it threw the error.

Commit message: chlog

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3d84bb7067f53..0323e81349aec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed RMSLE metric ([#3188](https://github.com/PyTorchLightning/pytorch-lightning/pull/3188))
 
+- Fixed ONNX model save on GPU ([#3145](https://github.com/PyTorchLightning/pytorch-lightning/pull/3145))
+
 ## [0.9.0] - YYYY-MM-DD
 
 ### Added
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index a726cf31fe365..61280dc2d776f 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1716,11 +1716,13 @@ def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwarg
         elif self.example_input_array is not None:
             input_data = self.example_input_array
         else:
-            raise ValueError('`input_sample` and `example_input_array` tensors are both missing.')
-
+            raise ValueError('Could not export to ONNX since neither `input_sample` nor'
+                             ' `model.example_input_array` attribute is set.')
+        input_data = input_data.to(self.device)
         if 'example_outputs' not in kwargs:
             self.eval()
-            kwargs['example_outputs'] = self(input_data)
+            with torch.no_grad():
+                kwargs['example_outputs'] = self(input_data)
 
         torch.onnx.export(self, input_data, file_path, **kwargs)
 
diff --git a/tests/models/test_onnx.py b/tests/models/test_onnx.py
index 278465941a043..4b9b81b1a6a1d 100644
--- a/tests/models/test_onnx.py
+++ b/tests/models/test_onnx.py
@@ -17,7 +17,21 @@ def test_model_saves_with_input_sample(tmpdir):
     trainer = Trainer(max_epochs=1)
     trainer.fit(model)
 
-    file_path = os.path.join(tmpdir, "model.onxx")
+    file_path = os.path.join(tmpdir, "model.onnx")
+    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
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
+def test_model_saves_on_gpu(tmpdir):
+    """Test that model saves on gpu"""
+    model = EvalModelTemplate()
+    trainer = Trainer(gpus=1, max_epochs=1)
+    trainer.fit(model)
+
+    file_path = os.path.join(tmpdir, "model.onnx")
     input_sample = torch.randn((1, 28 * 28))
     model.to_onnx(file_path, input_sample)
     assert os.path.isfile(file_path)
@@ -30,7 +44,7 @@ def test_model_saves_with_example_output(tmpdir):
     trainer = Trainer(max_epochs=1)
     trainer.fit(model)
 
-    file_path = os.path.join(tmpdir, "model.onxx")
+    file_path = os.path.join(tmpdir, "model.onnx")
     input_sample = torch.randn((1, 28 * 28))
     model.eval()
     example_outputs = model.forward(input_sample)
@@ -41,7 +55,7 @@ def test_model_saves_with_example_output(tmpdir):
 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")
+    file_path = os.path.join(tmpdir, "model.onnx")
     model.to_onnx(file_path)
     assert os.path.exists(file_path) is True
     assert os.path.getsize(file_path) > 3e+06
@@ -66,7 +80,7 @@ def test_model_saves_on_multi_gpu(tmpdir):
 
     tpipes.run_model_test(trainer_options, model)
 
-    file_path = os.path.join(tmpdir, "model.onxx")
+    file_path = os.path.join(tmpdir, "model.onnx")
     model.to_onnx(file_path)
     assert os.path.exists(file_path) is True
 
@@ -74,7 +88,7 @@ def test_model_saves_on_multi_gpu(tmpdir):
 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")
+    file_path = os.path.join(tmpdir, "model.onnx")
     model.to_onnx(file_path, verbose=True)
     captured = capsys.readouterr()
     assert "graph(%" in captured.out
@@ -84,8 +98,9 @@ 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'):
+    file_path = os.path.join(tmpdir, "model.onnx")
+    with pytest.raises(ValueError, match=r'Could not export to ONNX since neither `input_sample` nor'
+                                         r' `model.example_input_array` attribute is set.'):
         model.to_onnx(file_path)
 
 
@@ -99,7 +114,7 @@ def test_if_inference_output_is_valid(tmpdir):
     with torch.no_grad():
         torch_out = model(model.example_input_array)
 
-    file_path = os.path.join(tmpdir, "model.onxx")
+    file_path = os.path.join(tmpdir, "model.onnx")
     model.to_onnx(file_path, model.example_input_array, export_params=True)
 
     ort_session = onnxruntime.InferenceSession(file_path)


Code review comments:
diff_hunk: @@ -1716,11 +1716,13 @@ def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwarg
         elif self.example_input_array is not None:
             input_data = self.example_input_array
         else:
-            raise ValueError('`input_sample` and `example_input_array` tensors are both missing.')
-
+            raise ValueError('Could not export to ONNX since neither `input_sample` nor'
+                             ' `model.example_input_array` attribute is set.')
path: pytorch_lightning/core/lightning.py
text: @user1:
so how would you change it?
I see the condition above and it states as the error say - if `input_sample` is not tensor then it is not valid and `model.example_input_array is None` means it is missing...

@user2:
```suggestion
            if input_sample is not None:
                raise ValueError('`input_sample` should be a Tensor')
            else:
                raise ValueError('Could not export to ONNX since neither `input_sample` nor'
                                 ' `model.example_input_array` attribute is set.')
```
I don't know, maybe this?

@user1:
then lest be more spec with the first one 
```
raise ValueError(f'passed `input_sample` if {type(input_sample)} but it should be a Tensor')
```

@author:
@user2 good catch 👍 . I'll add the changes as @user1 suggested. :]