Repository: pytorch/pytorch

Issue:
[ONNX] onnx.export fails on frozen top-level ScriptedModule
### 🐛 Describe the bug

Here, onnx.export fails if called on a ScriptedModule that has been frozen:
```python
import torch

class Inner(torch.nn.Module):
    def forward(self, x):
        if x > 0 :  
            return x
        else:
            return x*x

class Outer(torch.nn.Module):
    def __init__(self):   
        super().__init__()
        i = Inner()
        self.inner = torch.jit.script(i)

    def forward(self, x):   
        return self.inner(x)

x = torch.zeros(1)
o=Outer()
o.eval()
m = torch.jit.trace_module(o, { 'forward' : (x)})
# borisf: passes if you comment this line out                                                                                                                                                        
m = torch.jit.optimize_for_inference(torch.jit.freeze(m))

torch.onnx.export(m, (x,), 'test.onnx')
```

Result:
python3 not-part.py
Traceback (most recent call last):
  File "not-part.py", line 26, in <module>
    torch.onnx.export(m, (x,), 'test.onnx')
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 479, in export
    _export(
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 1390, in _export
    with exporter_context(model, training, verbose):
  File "/opt/conda/lib/python3.8/contextlib.py", line 113, in __enter__
    return next(self.gen)
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 175, in exporter_context
    with select_model_mode_for_export(
  File "/opt/conda/lib/python3.8/contextlib.py", line 113, in __enter__
    return next(self.gen)
  File "/opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py", line 93, in select_model_mode_for_export
    originally_training = model.training
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 785, in __getattr__
    return super(RecursiveScriptModule, self).__getattr__(attr)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 502, in __getattr__
    return super(ScriptModule, self).__getattr__(attr)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1261, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'RecursiveScriptModule' object has no attribute 'training'




### Versions

Collecting environment information...
PyTorch version: 1.13.0a0+d321be6
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.23.3
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.7.99
GPU models and configuration: GPU 0: NVIDIA RTX A5000
Nvidia driver version: 515.65.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.5.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] k2==1.19.dev20220911+cuda11.7.torch1.13.0a0
[pip3] numpy==1.22.4
[pip3] pytorch-lightning==1.7.5
[pip3] pytorch-quantization==2.1.2
[pip3] torch==1.13.0a0+d321be6
[pip3] torch-tensorrt==1.2.0a0
[pip3] torchaudio==0.13.0
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.14.0a0
[conda] k2                        1.19.dev20220911+cuda11.7.torch1.13.0a0          pypi_0    pypi
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-include               2020.4             h726a3e6_304    conda-forge
[conda] numpy                     1.22.4           py38h99721a1_0    conda-forge
[conda] pytorch-lightning         1.7.5                    pypi_0    pypi
[conda] pytorch-quantization      2.1.2                    pypi_0    pypi
[conda] torch                     1.13.0a0+d321be6          pypi_0    pypi
[conda] torch-tensorrt            1.2.0a0                  pypi_0    pypi
[conda] torchaudio                0.13.0                   pypi_0    pypi
[conda] torchmetrics              0.9.3                    pypi_0    pypi
[conda] torchvision               0.14.0a0                 pypi_0    pypi



Discussion:
Possible fix in select_model_mode_for_export: 
```diff
root@30f44d0dcd3f:/git/radtts# diff -ur /opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py.orig /opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py
--- /opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py.orig     2022-10-05 21:20:37.425721244 +0000
+++ /opt/conda/lib/python3.8/site-packages/torch/onnx/utils.py  2022-10-05 21:37:07.644794060 +0000
@@ -89,7 +89,7 @@
         )
     originally_training: bool = False

-    if not isinstance(model, torch.jit.ScriptFunction):
+    if hasattr(model, 'training'):
         originally_training = model.training

         # ONNX opset 12 has better support for training amenable models, with updated
@@ -118,10 +118,7 @@
     try:
         yield
     finally:
-        if not (
-            isinstance(model, torch.jit.ScriptFunction)
-            or mode == _C_onnx.TrainingMode.PRESERVE
-        ):
+        if hasattr(model, 'training') and not mode == _C_onnx.TrainingMode.PRESERVE:
             model.train(originally_training)
```
Thanks for reporting this issue!
Let me show a workaround for guys who have the same issue. 

You may get an error `'RecursiveScriptModule' object has no attribute 'training'` when you try to convert your TorchScript into ONNX like the following:
```python
import torch

# loading your TorchScript
model = torch.jit.load("model.pt")

# converting the model to ONNX
dummy_input = ...
torch.onnx.export(model, dummy_input, "model.onnx")
```

You can avoid the error by adding the field explicitly:
```python
import torch

# loading your TorchScript
model = torch.jit.load("model.pt")

# ADD THIS LINE
model.training = False

# converting the model to ONNX
dummy_input = ...
torch.onnx.export(model, dummy_input, "model.onnx")
```


> Possible fix in select_model_mode_for_export:
> 

Thanks for the suggestion!

Resolving patch:
diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py
--- a/torch/onnx/utils.py
+++ b/torch/onnx/utils.py
@@ -90,7 +90,7 @@ def select_model_mode_for_export(model, mode: _C_onnx.TrainingMode):
         )
     originally_training: bool = False
 
-    if not isinstance(model, torch.jit.ScriptFunction):
+    if hasattr(model, "training"):
         originally_training = model.training
 
         # ONNX opset 12 has better support for training amenable models, with updated
@@ -119,10 +119,7 @@ def select_model_mode_for_export(model, mode: _C_onnx.TrainingMode):
     try:
         yield
     finally:
-        if not (
-            isinstance(model, torch.jit.ScriptFunction)
-            or mode == _C_onnx.TrainingMode.PRESERVE
-        ):
+        if hasattr(model, "training") and not mode == _C_onnx.TrainingMode.PRESERVE:
             model.train(originally_training)
 
 


Test patch:
diff --git a/test/onnx/test_utility_funs.py b/test/onnx/test_utility_funs.py
--- a/test/onnx/test_utility_funs.py
+++ b/test/onnx/test_utility_funs.py
@@ -3,9 +3,11 @@
 import copy
 import functools
 import io
+import warnings
 from typing import Callable
 
 import onnx
+import parameterized
 
 import torch
 import torch.onnx
@@ -17,7 +19,7 @@
     skipIfUnsupportedMaxOpsetVersion,
     skipIfUnsupportedMinOpsetVersion,
 )
-from torch.onnx import OperatorExportTypes, TrainingMode, utils
+from torch.onnx import _constants, OperatorExportTypes, TrainingMode, utils
 from torch.onnx._globals import GLOBALS
 from torch.onnx.symbolic_helper import _unpack_list, parse_args
 from torch.testing._internal import common_utils
@@ -128,8 +130,15 @@ def test_it_returns_empty_list_when_all_ops_convertible(
         self.assertEqual(unconvertible_ops, [])
 
 
-class TestUtilityFuns_opset9(_BaseTestCase):
-    opset_version = 9
+@parameterized.parameterized_class(
+    [
+        {"opset_version": opset}
+        for opset in range(_constants.ONNX_BASE_OPSET, _constants.ONNX_MAX_OPSET + 1)
+    ],
+    class_name_func=lambda cls, num, params_dict: f"{cls.__name__}_opset_{params_dict['opset_version']}",
+)
+class TestUtilityFuns(_BaseTestCase):
+    opset_version = None
 
     def test_is_in_onnx_export(self):
         test_self = self
@@ -148,8 +157,6 @@ def forward(self, x):
             self.assertFalse(torch.onnx.is_in_onnx_export())
 
     def test_validate_dynamic_axes_invalid_input_output_name(self):
-        import warnings
-
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
             utils._validate_dynamic_axes(
@@ -792,6 +799,31 @@ def forward(self, x):
         # verify that the model state is preserved
         self.assertEqual(model.training, old_state)
 
+    def test_export_does_not_fail_on_frozen_scripted_module(self):
+        class Inner(torch.nn.Module):
+            def forward(self, x):
+                if x > 0:
+                    return x
+                else:
+                    return x * x
+
+        class Outer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.inner = torch.jit.script(Inner())
+
+            def forward(self, x):
+                return self.inner(x)
+
+        x = torch.zeros(1)
+        # Freezing is only implemented in eval mode. So we need to call eval()
+        outer_module = Outer().eval()
+        module = torch.jit.trace_module(outer_module, {"forward": (x)})
+        # jit.freeze removes the training attribute in the module
+        module = torch.jit.freeze(module)
+
+        torch.onnx.export(module, (x,), io.BytesIO(), opset_version=self.opset_version)
+
     @skipIfUnsupportedMinOpsetVersion(15)
     def test_local_function(self):
         class N(torch.nn.Module):
@@ -1059,20 +1091,20 @@ def forward(self, x, y, z):
 
         model = M(3)
         expected_scope_names = {
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
             "torch.nn.modules.activation.GELU::gelu1",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
             "torch.nn.modules.activation.GELU::gelu2",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
             "torch.nn.modules.normalization.LayerNorm::lns.0",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
             "torch.nn.modules.normalization.LayerNorm::lns.1",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
             "torch.nn.modules.normalization.LayerNorm::lns.2",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::/"
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.N::relu/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::/"
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.N::relu/"
             "torch.nn.modules.activation.ReLU::relu",
-            "test_utility_funs.TestUtilityFuns_opset9.test_node_scope.<locals>.M::",
+            "test_utility_funs.TestUtilityFuns.test_node_scope.<locals>.M::",
         }
 
         graph, _, _ = self._model_to_graph(
@@ -1130,7 +1162,7 @@ def forward(self, x):
         #       so we expect 3 constants with different scopes. The 3 constants are for the 3 layers.
         #       If CSE in exporter is improved later, this test needs to be updated.
         #       It should expect 1 constant, with same scope as root.
-        scope_prefix = "test_utility_funs.TestUtilityFuns_opset9.test_scope_of_constants_when_combined_by_cse_pass.<locals>"
+        scope_prefix = "test_utility_funs.TestUtilityFuns.test_scope_of_constants_when_combined_by_cse_pass.<locals>"
         expected_root_scope_name = f"{scope_prefix}.N::"
         expected_layer_scope_name = f"{scope_prefix}.M::layers"
         expected_constant_scope_name = [
@@ -1180,7 +1212,7 @@ def forward(self, x):
         graph, _, _ = self._model_to_graph(
             N(), (torch.randn(2, 3)), input_names=[], dynamic_axes={}
         )
-        scope_prefix = "test_utility_funs.TestUtilityFuns_opset9.test_scope_of_nodes_when_combined_by_cse_pass.<locals>"
+        scope_prefix = "test_utility_funs.TestUtilityFuns.test_scope_of_nodes_when_combined_by_cse_pass.<locals>"
         expected_root_scope_name = f"{scope_prefix}.N::"
         expected_layer_scope_name = f"{scope_prefix}.M::layers"
         expected_add_scope_names = [
@@ -1884,29 +1916,5 @@ def forward(self, x):
         torch.onnx.unregister_custom_op_symbolic("::cat", _onnx_opset_version)
 
 
-class TestUtilityFuns_opset10(TestUtilityFuns_opset9):
-    opset_version = 10
-
-
-class TestUtilityFuns_opset11(TestUtilityFuns_opset9):
-    opset_version = 11
-
-
-class TestUtilityFuns_opset12(TestUtilityFuns_opset9):
-    opset_version = 12
-
-
-class TestUtilityFuns_opset13(TestUtilityFuns_opset9):
-    opset_version = 13
-
-
-class TestUtilityFuns_opset14(TestUtilityFuns_opset9):
-    opset_version = 14
-
-
-class TestUtilityFuns_opset15(TestUtilityFuns_opset9):
-    opset_version = 15
-
-
 if __name__ == "__main__":
     common_utils.run_tests()