Repository: Lightning-AI/pytorch-lightning

Issue:
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.



Discussion:

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
@@ -1716,11 +1716,16 @@ 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.')
-
+            if input_sample is not None:
+                raise ValueError(f'Received `input_sample` of type {type(input_sample)}. Expected type is `Tensor`')
+            else:
+                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)
 


Test patch:
diff --git a/tests/models/test_onnx.py b/tests/models/test_onnx.py
--- 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,11 +98,23 @@ 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)
 
 
+def test_error_if_input_sample_is_not_tensor(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.onnx")
+    input_sample = np.random.randn(1, 28 * 28)
+    with pytest.raises(ValueError, match=f'Received `input_sample` of type {type(input_sample)}. Expected type is '
+                                         f'`Tensor`'):
+        model.to_onnx(file_path, input_sample)
+
+
 def test_if_inference_output_is_valid(tmpdir):
     """Test that the output inferred from ONNX model is same as from PyTorch"""
     model = EvalModelTemplate()
@@ -99,7 +125,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)