[ONNX] Support opset 17 operators

[justinchuby]

Tracking issue on the onnx opset 17 support. Contributions welcome.

• Add version 17 as a known version [ONNX] Create empty opset 17 symbolic file #83287
• fft
  • ifft
  • fft_fft
  • fft_ifft
  • fft_rfft
  • fft_irfft
  • fft_hfft
  • fft_ihfft
  • fft_fft2
  • fft_ifft2
  • fft_rfft2
  • fft_irfft2
  • fft_hfft2
  • fft_ihfft2
  • fft_fftn
  • fft_ifftn
  • fft_rfftn
  • fft_irfftn
  • fft_hfftn
  • fft_ihfftn
  • fft_fftfreq
  • fft_rfftfreq
  • fft_fftshift
  • fft_ifftshift
• stft [ONNX] Opset 17 STFT support #83944
• istft
• bartlett_window
• blackman_window
• hamming_window
• hann_window
• kaiser_window
• LayerNormalization [ONNX] Add onnx::LayerNorm support for version 17 #84293

Closes #80834

[iver56]

Justin, you listed istft here, but I can't find istft in ONNX's list of operators. Is the idea to implement istft in terms of the available onnx operators?

Edit: This PR also lists ISTFT in its description, but there's no mention of ISTFT in the code changes. Now I'm confused. Does ONNX opset 17 support ISTFT or not?

[justinchuby]

Hi @iver56, I can confirm that ISTFT is not in opset 17. However onnx is open to including it as a new op (please feel free to open an issue in the onnx/onnx repo if not already existed). We plan to support exporting istft from pytorch in spite of the missing ISTFI op.

[stonelazy]

If I may add, even though ISTFT is not available, both forward and inverse operations for DFT are supported by ONNX. Reference This should help users implement ISTFT on their own.
Would it be possible to include support for FFT as well ?

[justinchuby]

@stonelazy i think so. Do you have a list of related torch functions you think could be fitting?

[stonelazy]

@stonelazy i think so. Do you have a list of related torch functions you think could be fitting?

I hope am not wrong, below are the torch's equivalent in ONNX.
Torch's FFT's equivalent is ONNX's DFT
Torch's IFFT is same ONNX's DFT but with inverse option selected.
Torch's Melscale equivalent is ONNX's Melweight.

[jonas-doevenspeck]

SInce LayerNorm is also part of opset 17, does it make sense to include it as well in this list?
https://github.com/onnx/onnx/blob/main/docs/Operators.md#layernormalization

[FrankFundel]

Will opset17 be available, only if all operators are implemented?

[justinchuby]

opset17 is enabled in torch-nightly. support for new ops will be added gradually.

[FrankFundel]

Thanks :) Hoping for stft and window functions to come soon.

[mravanelli]

That's a very important feature for SpeechBrain as well. The feature extraction part relies on fft/stft and we cannot currently fully export/import it in ONNX (@fpaissan).

[frankiedrake]

opset17 is enabled in torch-nightly. support for new ops will be added gradually.

Hello, I've installed torch from nightly releases (Version: 1.14.0.dev20221017+cu117) but I'm still getting an error
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator '::stft' to ONNX opset version 17 is not supported
Is that me installed the wrong version or this feature isn't released yet?

[justinchuby]

Hi! Even though we can export to the opset17 format, support for stft is not yet added unfortunately.

[david-macleod]

Are there any plans to implement fft_fft / fft_rfft or would you accept pull requests?

[justinchuby]

We will accept pull requests for sure!

[david-macleod]

Looks like it is blocked on support for complex types in onnx export 😔

[Alexey-Kamenev]

If you only need an ONNX support for real-valued FFT ops (e.g. rfft/irfft, rfft2/irfft2 and so on), you can use ONNX Contrib ops, such as Rfft and Irfft. You can take a look at my project which provides ONNX export + TensorRT plugins for models using FFT ops. That projects provides only a simple example of how Contrib ops can be used in ONNX export. For a complete implementation, refer to my code in NVIDIA Modulus project.

Note that if you plan to use ONNX models with those ops in runtime, you might need to use a custom build of ORT due to the bug in IRFFT. The alternative is to use TRT plugins, that is, do PyTorch -> ONNX -> TRT conversion (TRT plugins do not use ORT, so are not affected by the bug).

[sammlapp]

If you only need an ONNX support for real-valued FFT ops (e.g. rfft/irfft, rfft2/irfft2 and so on), you can use ONNX Contrib ops, such as Rfft and Irfft. You can take a look at my project which provides ONNX export + TensorRT plugins for models using FFT ops. That projects provides only a simple example of how Contrib ops can be used in ONNX export. For a complete implementation, refer to my code in NVIDIA Modulus project.

Note that if you plan to use ONNX models with those ops in runtime, you might need to use a custom build of ORT due to the bug in IRFFT. The alternative is to use TRT plugins, that is, do PyTorch -> ONNX -> TRT conversion (TRT plugins do not use ORT, so are not affected by the bug).

@Alexey-Kamenev this sounds like it could potentially be extremely useful, but I don't really understand what I would need to do. For instance, could I create a PyTorch model that creates spectrograms from audio, and export this model to ONNX? If so, could you describe the steps necessary?

[Alexey-Kamenev]

@sammlapp

For instance, could I create a PyTorch model that creates spectrograms from audio, and export this model to ONNX?

Yes, you could do that, if in your PyTorch code you use something like torch.fft.rfft.
There is also an example of using ONNX-friendly FFT functions to implement AFNO (Adaptive Fourier Neural Operator): code, paper.

If so, could you describe the steps necessary?

Unfortunately, due to lack of support of complex type in ONNX exporter, the process is not exactly trivial - you would need to use some sort of a wrapper which would use standard PyTorch FFT functions when not in ONNX export mode and custom symbolic functions which emit proper ONNX operators (e.g. Rfft or Irfft) during ONNX export. Also, you would need to change your code to avoid using complex-typed tensors, essentially treating all complex tensors as if after torch.view_as_real.

We already did all of that in Modulus (see the AFNO example I mentioned above), but if you don't want to take the dependency on Modulus, you can implement a simpler version using the code by checking the tests.

I know, it does not sound particularly simple, but that's the consequence of not having complex type support in ONNX exporter. Otherwise, it would have been as simple as something like calling register_custom_op_symbolic.

[justinchuby]

More support on the way in our new torch.onnx.dynamo_export exporter. cc @titaiwangms

[shingjan]

There is a aten:l1_loss not supported error for model Super_SloMo in torchbench. Issue reported here. Wonder if this op will be supported in the future? Although it is odd that a training-related op is included for inference.

[justinchuby]

l1_loss will be supported

[TechInterMezzo]

What is the current workaround to export models with mel spectrograms? And is support for stft the only thing that is missing to make it work?

[justinchuby]

STFT is supported by the torch.onnx.dynamo_export export API. I have not looked at melspectrogram. Supposedly ONNX has a related op that we have not made use of.

[justinchuby]

It should now be supported by torch.onnx.dynamo_export. Note that the dynamo exporter for ONNX is in beta.

[sammlapp]

@justinchuby does this mean that there's now a simpler method for exporting pytorch models containing fft than the approach described by @Alexey-Kamenev above? Or is that approach still required?
Thanks

[justinchuby]

Yes. The new torch.onnx.dynamo_export should allow a simpler experience for exporting FFT operators. The approach described by @Alexey-Kamenev will be required for the torch.onnx.export API.

[shanecarroll-smarsh]

Can anyone produce a working example where torch.onnx.dynamo_export successfully exports a torch.stft op?

Here is a simple MWE, with a setup common to audio signal processing models:

import torch


class STFTModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._window = torch.hann_window(window_length=320)

    def forward(self, signals: torch.Tensor) -> torch.Tensor:
        x = signals.stft(
            n_fft=512,
            hop_length=160,
            win_length=320,
            return_complex=True,  # doesn't affect errors
            window=self._window,
            pad_mode="constant",  # aten.reflection_pad1d unsupported op
        )
        return x


m = STFTModel()

# Shape [B, T] audio signals
input_signals = torch.randn([2, 16000])

args = (input_signals,)
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
torch.onnx.dynamo_export(
    m,
    *args,
    export_options=export_options,
)

Here are the short versions of error messages:

Without dynamic shapes (not useful to anyone using stft):

torch.onnx._internal.diagnostics.infra.context.RuntimeErrorWithDiagnostic: Unsupported FX nodes: {'call_function': ['aten.transpose.int']}. 

With dynamic shapes (as the example shows):

torch._dynamo.exc.Unsupported: unsupported operator: aten._fft_r2c.default

Exporting within the context of torch.inference_mode(), output is slightly different (prims vs. aten):

torch._dynamo.exc.Unsupported: unsupported operator: prims.fft_r2c.default

Relevant context (should be the latest in everything):

$ pip freeze | egrep '(torch|onnx)'
onnx==1.15.0
onnxscript==0.1.0.dev20231106
pytorch-triton==2.1.0+6e4932cda8
torch==2.2.0.dev20231106+cu121
torchaudio==2.2.0.dev20231106+cu121
torchvision==0.17.0.dev20231106+cu121

[justinchuby]

Can anyone produce a working example where torch.onnx.dynamo_export successfully exports a torch.stft op?

Here is a simple MWE, with a setup common to audio signal processing models:

import torch


class STFTModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._window = torch.hann_window(window_length=320)

    def forward(self, signals: torch.Tensor) -> torch.Tensor:
        x = signals.stft(
            n_fft=512,
            hop_length=160,
            win_length=320,
            return_complex=True,  # doesn't affect errors
            window=self._window,
            pad_mode="constant",  # aten.reflection_pad1d unsupported op
        )
        return x


m = STFTModel()

# Shape [B, T] audio signals
input_signals = torch.randn([2, 16000])

args = (input_signals,)
export_options = torch.onnx.ExportOptions(dynamic_shapes=True)
torch.onnx.dynamo_export(
    m,
    *args,
    export_options=export_options,
)

Here are the short versions of error messages:

Without dynamic shapes (not useful to anyone using stft):

torch.onnx._internal.diagnostics.infra.context.RuntimeErrorWithDiagnostic: Unsupported FX nodes: {'call_function': ['aten.transpose.int']}. 

With dynamic shapes (as the example shows):

torch._dynamo.exc.Unsupported: unsupported operator: aten._fft_r2c.default

Exporting within the context of torch.inference_mode(), output is slightly different (prims vs. aten):

torch._dynamo.exc.Unsupported: unsupported operator: prims.fft_r2c.default

Thanks for sharing the info. I will reproduce this and get back to you.

[justinchuby]

We will continue in #113067 @shanecarroll-smarsh