torchscript support

[seyong92]

Hi, thank you for such an amazing project!

I want to suggest the support of torchscript for production use of nnAudio.

Currently, it does not work with the command script_model = torch.jit.script(model) with following errors.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[6], line 1
----> 1 sm = torch.jit.script(model)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_script.py:1324, in script(obj, optimize, _frames_up, _rcb, example_inputs)
   1322 if isinstance(obj, torch.nn.Module):
   1323     obj = call_prepare_scriptable_func(obj)
-> 1324     return torch.jit._recursive.create_script_module(
   1325         obj, torch.jit._recursive.infer_methods_to_compile
   1326     )
   1327 else:
   1328     obj = obj.__prepare_scriptable__() if hasattr(obj, "__prepare_scriptable__") else obj  # type: ignore[operator]

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:559, in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    557 if not is_tracing:
    558     AttributeTypeIsSupportedChecker().check(nn_module)
--> 559 return create_script_module_impl(nn_module, concrete_type, stubs_fn)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:632, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    629     script_module._concrete_type = concrete_type
    631 # Actually create the ScriptModule, initializing it with the function we just defined
--> 632 script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    634 # Compile methods if necessary
    635 if concrete_type not in concrete_type_store.methods_compiled:

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_script.py:639, in RecursiveScriptModule._construct(cpp_module, init_fn)
    626 """
    627 Construct a RecursiveScriptModule that's ready for use. PyTorch
    628 code should use this to construct a RecursiveScriptModule instead
   (...)
    636     init_fn:  Lambda that initializes the RecursiveScriptModule passed to it.
    637 """
    638 script_module = RecursiveScriptModule(cpp_module)
--> 639 init_fn(script_module)
    641 # Finalize the ScriptModule: replace the nn.Module state with our
    642 # custom implementations and flip the _initializing bit.
    643 RecursiveScriptModule._finalize_scriptmodule(script_module)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:608, in create_script_module_impl.<locals>.init_fn(script_module)
    605     scripted = orig_value
    606 else:
    607     # always reuse the provided stubs_fn to infer the methods to compile
--> 608     scripted = create_script_module_impl(
    609         orig_value, sub_concrete_type, stubs_fn
    610     )
    612 cpp_module.setattr(name, scripted)
    613 script_module._modules[name] = scripted

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:632, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    629     script_module._concrete_type = concrete_type
    631 # Actually create the ScriptModule, initializing it with the function we just defined
--> 632 script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    634 # Compile methods if necessary
    635 if concrete_type not in concrete_type_store.methods_compiled:

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_script.py:639, in RecursiveScriptModule._construct(cpp_module, init_fn)
    626 """
    627 Construct a RecursiveScriptModule that's ready for use. PyTorch
    628 code should use this to construct a RecursiveScriptModule instead
   (...)
    636     init_fn:  Lambda that initializes the RecursiveScriptModule passed to it.
    637 """
    638 script_module = RecursiveScriptModule(cpp_module)
--> 639 init_fn(script_module)
    641 # Finalize the ScriptModule: replace the nn.Module state with our
    642 # custom implementations and flip the _initializing bit.
    643 RecursiveScriptModule._finalize_scriptmodule(script_module)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:608, in create_script_module_impl.<locals>.init_fn(script_module)
    605     scripted = orig_value
    606 else:
    607     # always reuse the provided stubs_fn to infer the methods to compile
--> 608     scripted = create_script_module_impl(
    609         orig_value, sub_concrete_type, stubs_fn
    610     )
    612 cpp_module.setattr(name, scripted)
    613 script_module._modules[name] = scripted

    [... skipping similar frames: RecursiveScriptModule._construct at line 639 (1 times), create_script_module_impl at line 632 (1 times), create_script_module_impl.<locals>.init_fn at line 608 (1 times)]

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:632, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    629     script_module._concrete_type = concrete_type
    631 # Actually create the ScriptModule, initializing it with the function we just defined
--> 632 script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
    634 # Compile methods if necessary
    635 if concrete_type not in concrete_type_store.methods_compiled:

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_script.py:639, in RecursiveScriptModule._construct(cpp_module, init_fn)
    626 """
    627 Construct a RecursiveScriptModule that's ready for use. PyTorch
    628 code should use this to construct a RecursiveScriptModule instead
   (...)
    636     init_fn:  Lambda that initializes the RecursiveScriptModule passed to it.
    637 """
    638 script_module = RecursiveScriptModule(cpp_module)
--> 639 init_fn(script_module)
    641 # Finalize the ScriptModule: replace the nn.Module state with our
    642 # custom implementations and flip the _initializing bit.
    643 RecursiveScriptModule._finalize_scriptmodule(script_module)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:608, in create_script_module_impl.<locals>.init_fn(script_module)
    605     scripted = orig_value
    606 else:
    607     # always reuse the provided stubs_fn to infer the methods to compile
--> 608     scripted = create_script_module_impl(
    609         orig_value, sub_concrete_type, stubs_fn
    610     )
    612 cpp_module.setattr(name, scripted)
    613 script_module._modules[name] = scripted

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:636, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    634 # Compile methods if necessary
    635 if concrete_type not in concrete_type_store.methods_compiled:
--> 636     create_methods_and_properties_from_stubs(
    637         concrete_type, method_stubs, property_stubs
    638     )
    639     # Create hooks after methods to ensure no name collisions between hooks and methods.
    640     # If done before, hooks can overshadow methods that aren't exported.
    641     create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)

File /mypath/.venv/lib/python3.11/site-packages/torch/jit/_recursive.py:469, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    466 property_defs = [p.def_ for p in property_stubs]
    467 property_rcbs = [p.resolution_callback for p in property_stubs]
--> 469 concrete_type._create_methods_and_properties(
    470     property_defs, property_rcbs, method_defs, method_rcbs, method_defaults
    471 )

RuntimeError: 
Tried to set nonexistent attribute: num_samples. Did you forget to initialize it in __init__()?:
  File "/mypath/.venv/lib/python3.11/site-packages/nnAudio/features/stft.py", line 275
        """
        output_format = output_format or self.output_format
        self.num_samples = x.shape[-1]
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    
        x = broadcast_dim(x)

Thank you :)

[KinWaiCheuk]

Hi, thank you for your comments.
Recently, I'm a little bit busy with other commitments.

I will have a look as soon as I'm free, but I couldn't gurantee anything at the moment.
Sorry about that.

[migperfer]

Hi! Which audio transformation were you using? Would it be possible for you to share a code snippet of the model?

[dorienh]

TorchScript compatibility has been fixed in nnAudio2, the modernised successor to this repo: https://github.com/AMAAI-Lab/nnAudio2

The compilation failures in STFT and iSTFT were caused by dynamic state mutation and module construction in scripted code paths — these have been resolved. You can install it via pip install nnaudio2.