Exporting onnx and jit model at the same time

yunusemre: <p>Hi, I am trying to export a torch model to onnx and jit at the same time for a library (<a href="https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/pruned_transducer_stateless7/export.py" rel="noopener nofollow ugc">Icefall</a>). Some parts of models have to be changed to export. I am seperating these parts with  <code>torch.jit.is_scripting()</code> and <code>torch.onnx.is_in_onnx_export()</code>. I can export onnx model, but cannot export jit model. A simple example:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 2)
        self.fc2 = nn.Linear(10, 2)

    def forward(self, x):
        if torch.jit.is_scripting():
            x = self.fc1(x)
    
        if torch.onnx.is_in_onnx_export():
            x = self.fc2(x)
    
        return x

net = Net()
sample = torch.randn(1, 10)

torch.onnx.export(
    net,
    (sample),
    "model.onnx",
    verbose=False,
    opset_version=13,
    input_names=["x"],
    output_names=["out"],
)

net = torch.jit.script(net)
net.save(str("model.pt"))
</code></pre>
<p>Onnx export is done successfully, but  jit scripting raise error:</p>
<pre><code class="lang-auto">UnsupportedNodeError                      Traceback (most recent call last)
Cell In [9], line 1
----&gt; 1 net = torch.jit.script(net)
      2 net.save(str("model.pt"))

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_script.py:1257, in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1255 if isinstance(obj, torch.nn.Module):
   1256     obj = call_prepare_scriptable_func(obj)
-&gt; 1257     return torch.jit._recursive.create_script_module(
   1258         obj, torch.jit._recursive.infer_methods_to_compile
   1259     )
   1261 if isinstance(obj, dict):
   1262     return create_script_dict(obj)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_recursive.py:451, in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    449 if not is_tracing:
    450     AttributeTypeIsSupportedChecker().check(nn_module)
--&gt; 451 return create_script_module_impl(nn_module, concrete_type, stubs_fn)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_recursive.py:517, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    515 # Compile methods if necessary
    516 if concrete_type not in concrete_type_store.methods_compiled:
--&gt; 517     create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    518     # Create hooks after methods to ensure no name collisions between hooks and methods.
    519     # If done before, hooks can overshadow methods that aren't exported.
    520     create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_recursive.py:368, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    365 property_defs = [p.def_ for p in property_stubs]
    366 property_rcbs = [p.resolution_callback for p in property_stubs]
--&gt; 368 concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_recursive.py:838, in try_compile_fn(fn, loc)
    834 # We don't have the actual scope where the function was defined, but we can
    835 # extract the necessary info from the closed over variables on the function
    836 # object
    837 rcb = _jit_internal.createResolutionCallbackFromClosure(fn)
--&gt; 838 return torch.jit.script(fn, _rcb=rcb)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/_script.py:1307, in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1305 if maybe_already_compiled_fn:
   1306     return maybe_already_compiled_fn
-&gt; 1307 ast = get_jit_def(obj, obj.__name__)
   1308 if _rcb is None:
   1309     _rcb = _jit_internal.createResolutionCallbackFromClosure(obj)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/frontend.py:264, in get_jit_def(fn, def_name, self_name, is_classmethod)
    261     qualname = get_qualified_name(fn)
    262     pdt_arg_types = type_trace_db.get_args_types(qualname)
--&gt; 264 return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types)

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/frontend.py:315, in build_def(ctx, py_def, type_line, def_name, self_name, pdt_arg_types)
    310     type_comment_decl = torch._C.parse_type_comment(type_line)
    311     decl = torch._C.merge_type_from_type_comment(decl, type_comment_decl, is_method)
    313 return Def(Ident(r, def_name),
    314            decl,
--&gt; 315            build_stmts(ctx, body))

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/frontend.py:137, in build_stmts(ctx, stmts)
    136 def build_stmts(ctx, stmts):
--&gt; 137     stmts = [build_stmt(ctx, s) for s in stmts]
    138     return list(filter(None, stmts))

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/frontend.py:137, in &lt;listcomp&gt;(.0)
    136 def build_stmts(ctx, stmts):
--&gt; 137     stmts = [build_stmt(ctx, s) for s in stmts]
    138     return list(filter(None, stmts))

File ~/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/jit/frontend.py:286, in Builder.__call__(self, ctx, node)
    284 method = getattr(self, 'build_' + node.__class__.__name__, None)
    285 if method is None:
--&gt; 286     raise UnsupportedNodeError(ctx, node)
    287 return method(ctx, node)

UnsupportedNodeError: import statements aren't supported:
  File "/home/yunusemre.ozkose/anaconda3/envs/k2sh/lib/python3.8/site-packages/torch/onnx/__init__.py", line 386
    """

    from torch.onnx import utils
    ~~~~ &lt;--- HERE
    return utils.is_in_onnx_export()
</code></pre>
<p>How can I solve this?</p>
<p>My env:</p>
<pre><code class="lang-auto">Versions of relevant libraries:
[pip3] k2==1.21.dev20221103+cuda11.3.torch1.10.0
[pip3] numpy==1.22.4
[pip3] torch==1.10.0
[pip3] torchaudio==0.10.0
[pip3] torchvision==0.11.0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1              h9edb442_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] k2                        1.21.dev20221103+cuda11.3.torch1.10.0          pypi_0    pypi
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py38h95df7f1_0    conda-forge
[conda] mkl_fft                   1.3.1            py38h8666266_1    conda-forge
[conda] mkl_random                1.2.2            py38h1abd341_0    conda-forge
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] pytorch                   1.10.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.10.0               py38_cu113    pytorch
[conda] torchvision               0.11.0               py38_cu113    pytorch
</code></pre>