[ONNX] STFT Support #92087

urinieto · 2023-01-12T17:49:38Z

This PR addresses issue #81075, making torch.stft compatible with ONNX Opset 17's STFT operator.

The conversion works for most of torch.stft functionality:

Batched or unbatched inputs
Normalization
Pre-computed windows
Rectangular windows
One-sided returns
Window centering (implicitly supported)

What is currently not supported is complex types, due to the lack of conversion functionality between PyTorch and ONNX (#86746).

Regardless, this is easy to bypass by setting return_complex=False when using torch.stft.

Note that there is already a draft PR to address this (#83944), but it is currently closed and it only partially addresses the conversion (i.e., most of torch.stft functionality is lacking, and unit tests are missing).

linux-foundation-easycla · 2023-01-12T17:49:42Z

The committers listed above are authorized under a signed CLA.

✅ login: urinieto / name: Oriol Nieto (a95ffd7)

pytorch-bot · 2023-01-12T17:49:42Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92087

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 7 Failures

As of commit b88508c:

NEW FAILURES - The following jobs have failed:

linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 1, 4, linux.g5.4xlarge.nvidia.gpu) (gh)
linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 2, 4, linux.g5.4xlarge.nvidia.gpu) (gh)
linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 3, 4, linux.g5.4xlarge.nvidia.gpu) (gh)
linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (slow, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
win-vs2019-cpu-py3 / test (default, 1, 2, windows.4xlarge) (gh)
win-vs2019-cpu-py3 / test (default, 2, 2, windows.4xlarge) (gh)
win-vs2019-cpu-py3 / test (functorch, 1, 1, windows.4xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

justinchuby · 2023-01-13T02:21:26Z

Thanks for your contribution! Could you sign the CLA following instructions by the bot comment above, and fix lint issues by running lintrunner -a -m master? I will add comments after the CLA is signed.

https://github.com/pytorch/pytorch/wiki/lintrunner

urinieto · 2023-01-13T07:22:13Z

Thanks for reviewing this, @justinchuby! I just signed the CLA, and associated my last commits with the right email address. Moreover, I have passed the linter and fixed a problem with the unit tests regarding the return_complex parameter. Should be good to be reviewed :)

justinchuby

Looks like there's still errors in the CLA. Squashing the commits may work?

test/onnx/test_operators.py

torch/onnx/symbolic_opset17.py

urinieto · 2023-01-13T22:28:11Z

Apologies, I realized that the CLA should be signed under a "company contribution", since most of this work was done using my company's resources (Adobe). This contribution was internally approved before I submitted this PR, so I might need a working day or two to have someone from the Open Source Office sign the CLA. Will report asap.

milesial · 2023-01-15T19:49:28Z

Not sure if related to this PR or the underlying onnxruntime, but there seems to be a off-by-one error somewhere when using graph optimization with STFT.

class Test(nn.Module):
    def forward(self, audio):
        stft = torch.stft(audio, 400, 160, return_complex=False)
        return stft

m = Test()
print('torch', m(torch.randn(1, 16000 * 30)).shape)

torch.onnx.export(m,
                  torch.randn(1, 16000 * 30),
                  'test.onnx',
                  export_params=True,
                  input_names=['in'],
                  output_names=['out'],
                  opset_version=17, verbose=False)

sess_options = rt.SessionOptions()
sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_DISABLE_ALL
sess_options.optimized_model_filepath = 'test.onnx'
session = rt.InferenceSession('test.onnx', sess_options)
print('onnx unoptimized', session.run(None, {'in': torch.randn(1, 16000 * 30).numpy()})[0].shape)

sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_BASIC
sess_options.optimized_model_filepath = 'test.onnx'
session = rt.InferenceSession('test.onnx', sess_options)
print('onnx optimized', session.run(None, {'in': torch.randn(1, 16000 * 30).numpy()})[0].shape)

torch torch.Size([1, 201, 4801, 2])
onnx unoptimized (1, 201, 4801, 2)
onnx optimized (1, 201, 4801, 2)

2023-01-15 19:38:25.256968539 [W:onnxruntime:, execution_frame.cc:828 VerifyOutputSizes] Expected shape from model of {1,201,4802,2} does not match actual shape of {1,201,4801,2} for output out

The output shape matches, but the expected shape is different, and so when we add nodes on top of that it causes shape mismatch errors.

peterbell10 · 2023-01-16T15:50:09Z

torch/onnx/symbolic_opset17.py

+    # Checks
+    assert (
+        return_complex is None or not return_complex
+    ), "STFT does not currently support complex types"


return_complex=False is deprecated so I don't like the idea of forcing people to use it. I see that complex64 and complex128 are mentioned as types in this document, so is it possible to implement view_as_real/view_as_complex in onnx?

I like the idea of having view_as_real and view_as_complex ops in onnx. Would you be willing to open an issue for a new operator? https://github.com/onnx/onnx/issues/new?assignees=&labels=operator&template=operator.md&title=

@peterbell10 I also don't like the idea of forcing return_complex=False, but it is an easy workaround if we don't want to wait to have complex conversion support on ONNX, which seems like it might take quite some time (see below).

@justinchuby Seems like view_as_complex was already requested here: #49793 Unfortunately, there was not enough interest for this when it was reported. I just commented on the issue, to try to bring it back.

Am I right in saying onnx Cast doesn't support complex types? If it did then you could do something like:

def _as_real(z): return float(z), float(-1j * z) def _as_complex(real, imag): return complex(real) + 1j * complex(imag)

That is correct: ONNX doesn't currently support casting for complex types (https://github.com/onnx/onnx/blob/main/docs/Operators.md#Cast).

If it did, we could simply use Cast on the result of STFT to make sure the result is returned as complex if return_complex=True.

That being said, would it be possible to obtain the output of STFT, put it on a PyTorch tensor, and convert it to complex using the _as_complex function above, and then put it back into the ONNX Graph?

Sorry, just realized that that's not gonna work, because Cast from/to complex is not supported, which is exactly what @peterbell10 said (😅).

Yeah, otherwise there will be a graph break in onnx.

onnx/onnx#4785

urinieto · 2023-01-16T21:30:43Z

Not sure if related to this PR or the underlying onnxruntime, but there seems to be a off-by-one error somewhere when using graph optimization with STFT.

@milesial I noticed that as well. The fact that it actually returns the correct shape in all the tests I wrote makes me think it might be a bug on the the onnxruntime API (or the ONNX STFT definition?). I couldn't find an issue reported on the onnxruntime repo.

I was thinking of reporting it once this PR is merged, so that it might be easier to reproduce. Or if you wanna report it now, by all means do :)

milesial · 2023-01-16T23:25:03Z

Got it, good you caught it too, I opened microsoft/onnxruntime#14316 since I don't know if this MR is going to be merged soon

urinieto · 2023-02-09T19:35:00Z

The CLA is finally signed, apologies it took so long. I'll work on the comments and feedback asap.

test/onnx/expect/TestOperators.test_stft_default.expect

torch/onnx/symbolic_opset17.py

test/onnx/test_pytorch_onnx_onnxruntime.py

justinchuby · 2023-02-14T20:15:23Z

I notice the CLA bot is still having errors. Perhaps squash or rebase some commits?

justinchuby · 2023-02-15T00:13:06Z

There are some build errors I believe happened because this branch is too old. Could you rebase with master?

urinieto · 2023-03-01T02:28:42Z

You may run the test with pytest test/onnx/test_op_consistency.py -v -s -k stft

Hey @justinchuby , I fixed the problem regarding the type mismatch here.

The other problem (i.e., the small differences between the results), should be resolved setting the tolerance to around 1e-5 (like I did, e.g., here). I'm not familiar with how to change this tolerance in the test_op_consistency.py file, could you take care of it?

For reference, here's the last chunk of my output of pytest test/onnx/test_op_consistency.py -v -s -k stft:

FAILED test/onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset17CPU::test_output_match_stft_cpu_float32 - AssertionError: Tensor-likes are not close!
FAILED test/onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset17CPU::test_output_match_stft_cpu_float64 - AssertionError: Tensor-likes are not close!
FAILED test/onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset18CPU::test_output_match_stft_cpu_float32 - AssertionError: Tensor-likes are not close!
FAILED test/onnx/test_op_consistency.py::TestOnnxModelOutputConsistency_opset18CPU::test_output_match_stft_cpu_float64 - AssertionError: Tensor-likes are not close!

Thanks for your help!

justinchuby · 2023-03-02T17:37:02Z

Thanks! I will work on it next week when I am back

justinchuby

Thanks so much for creating this!

justinchuby · 2023-03-09T18:36:59Z

@pytorchbot merge -g

justinchuby · 2023-03-09T18:44:41Z

@pytorchbot merge -g

justinchuby · 2023-03-10T02:14:56Z

@pytorchbot merge

pytorchmergebot · 2023-03-10T02:16:54Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2023-03-10T02:17:00Z

Merge failed

Reason: 7 mandatory check(s) failed. The first few are:

pull / linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 1, 4, linux.g5.4xlarge.nvidia.gpu)
pull / linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 2, 4, linux.g5.4xlarge.nvidia.gpu)
pull / linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (default, 3, 4, linux.g5.4xlarge.nvidia.gpu)
pull / linux-bionic-cuda11.7-py3.10-gcc7-sm86 / test (slow, 2, 2, linux.g5.4xlarge.nvidia.gpu)
pull / win-vs2019-cpu-py3 / test (default, 1, 2, windows.4xlarge)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

justinchuby · 2023-03-10T02:19:13Z

@pytorchbot merge -f "unrelated cuda and other failures"

pytorchmergebot · 2023-03-10T02:20:54Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

This PR addresses issue [#81075](pytorch/pytorch#81075), making `torch.stft` compatible with ONNX Opset 17's STFT operator. The conversion works for _most_ of `torch.stft` functionality: - Batched or unbatched inputs - Normalization - Pre-computed windows - Rectangular windows - One-sided returns - Window centering (implicitly supported) What is currently _not_ supported is **complex types**, due to the lack of conversion functionality between PyTorch and ONNX (pytorch/pytorch#86746). Regardless, this is easy to bypass by setting `return_complex=False` when using `torch.stft`. Note that there is already a draft PR to address this (pytorch/pytorch#83944), but it is currently closed and it only partially addresses the conversion (i.e., most of `torch.stft` functionality is lacking, and unit tests are missing). Pull Request resolved: pytorch/pytorch#92087 Approved by: https://github.com/justinchuby

urinieto requested review from BowenBao and abock as code owners January 12, 2023 17:49

pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Jan 12, 2023

pytorchbot added the open source label Jan 12, 2023

BowenBao requested a review from justinchuby January 12, 2023 18:28

justinchuby self-assigned this Jan 12, 2023

drisspg added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jan 13, 2023

justinchuby added the module: onnx Related to torch.onnx label Jan 13, 2023

justinchuby reviewed Jan 13, 2023

View reviewed changes

divideconcept mentioned this pull request Jan 13, 2023

[New Operator Proposal] ISTFT operator onnx/onnx#4777

Open

divideconcept mentioned this pull request Jan 14, 2023

torch.istft should keep accepting real inputs in the future and not display a warning, because of ONNX #92204

Closed

peterbell10 reviewed Jan 16, 2023

View reviewed changes

urinieto mentioned this pull request Jan 16, 2023

[ONNX][Complex] Support view_as_complex #49793

Closed

milesial mentioned this pull request Jan 16, 2023

STFT op has the wrong expected shape microsoft/onnxruntime#14316

Open

justinchuby mentioned this pull request Jan 18, 2023

[New Operator Proposal] ViewAsReal/ViewAsComplex onnx/onnx#4785

Closed

ooe1123 mentioned this pull request Feb 12, 2023

ADD ReazonSpeech axinc-ai/ailia-models#1014

Open

justinchuby reviewed Feb 14, 2023

View reviewed changes

test/onnx/expect/TestOperators.test_stft_default.expect Outdated Show resolved Hide resolved

justinchuby reviewed Feb 14, 2023

View reviewed changes

torch/onnx/symbolic_opset17.py Outdated Show resolved Hide resolved

justinchuby reviewed Feb 14, 2023

View reviewed changes

test/onnx/test_pytorch_onnx_onnxruntime.py Outdated Show resolved Hide resolved

urinieto force-pushed the stft_onnx2 branch from f231cb4 to baa7129 Compare March 1, 2023 02:22

urinieto added 2 commits February 28, 2023 18:31

Addressing problem with missmatch types between input and window

e19ff4a

Copying changes in tests done in pytorch#95531

914314f

urinieto force-pushed the stft_onnx2 branch from baa7129 to 914314f Compare March 1, 2023 02:32

justinchuby mentioned this pull request Mar 9, 2023

[ONNX] Export logical_not #96315

Closed

justinchuby added 2 commits March 9, 2023 10:33

Adjust tolerance and format

e593ce8

Format

120a30b

justinchuby added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 9, 2023

justinchuby approved these changes Mar 9, 2023

View reviewed changes

Tolerance

5462f49

justinchuby removed the ciflow/trunk Trigger trunk jobs on your pull request label Mar 9, 2023

justinchuby added 2 commits March 9, 2023 13:41

Update test_op_consistency.py

2fdaaae

Update symbolic_opset17.py

b88508c

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 10, 2023

pytorchmergebot added the Merged label Mar 10, 2023

pytorchmergebot closed this in 5f89d14 Mar 10, 2023

This was referenced Aug 22, 2023

export model to onnx without FFT/IFFT facebookresearch/demucs#530

Open

onnx conversion facebookresearch/demucs#539

Open

garymm mentioned this pull request Feb 6, 2024

torch.stft export bug, using operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK, throws error: ONNX FATAL: [enforce fail at backend.cc:777] . Caffe2 only supports padding 2D Tensor, whereas padding is [0, 0, 160, 0, 0, 160, ] #18521

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[ONNX] STFT Support #92087

[ONNX] STFT Support #92087

urinieto commented Jan 12, 2023

linux-foundation-easycla bot commented Jan 12, 2023 •

edited

pytorch-bot bot commented Jan 12, 2023 •

edited

justinchuby commented Jan 13, 2023

urinieto commented Jan 13, 2023

justinchuby left a comment

urinieto commented Jan 13, 2023

milesial commented Jan 15, 2023

peterbell10 Jan 16, 2023

justinchuby Jan 16, 2023

urinieto Jan 16, 2023

peterbell10 Jan 17, 2023

urinieto Jan 17, 2023 •

edited

justinchuby Jan 17, 2023

justinchuby Mar 10, 2023

urinieto commented Jan 16, 2023

milesial commented Jan 16, 2023

urinieto commented Feb 9, 2023

justinchuby commented Feb 14, 2023

justinchuby commented Feb 15, 2023

urinieto commented Mar 1, 2023 •

edited

justinchuby commented Mar 2, 2023

justinchuby left a comment

justinchuby commented Mar 9, 2023

justinchuby commented Mar 9, 2023

justinchuby commented Mar 10, 2023

pytorchmergebot commented Mar 10, 2023

pytorchmergebot commented Mar 10, 2023

justinchuby commented Mar 10, 2023

pytorchmergebot commented Mar 10, 2023

[ONNX] STFT Support #92087

[ONNX] STFT Support #92087

Conversation

urinieto commented Jan 12, 2023

linux-foundation-easycla bot commented Jan 12, 2023 • edited

pytorch-bot bot commented Jan 12, 2023 • edited

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/92087

❌ 7 Failures

justinchuby commented Jan 13, 2023

urinieto commented Jan 13, 2023

justinchuby left a comment

Choose a reason for hiding this comment

urinieto commented Jan 13, 2023

milesial commented Jan 15, 2023

peterbell10 Jan 16, 2023

Choose a reason for hiding this comment

justinchuby Jan 16, 2023

Choose a reason for hiding this comment

urinieto Jan 16, 2023

Choose a reason for hiding this comment

peterbell10 Jan 17, 2023

Choose a reason for hiding this comment

urinieto Jan 17, 2023 • edited

Choose a reason for hiding this comment

justinchuby Jan 17, 2023

Choose a reason for hiding this comment

justinchuby Mar 10, 2023

Choose a reason for hiding this comment

urinieto commented Jan 16, 2023

milesial commented Jan 16, 2023

urinieto commented Feb 9, 2023

justinchuby commented Feb 14, 2023

justinchuby commented Feb 15, 2023

urinieto commented Mar 1, 2023 • edited

justinchuby commented Mar 2, 2023

justinchuby left a comment

Choose a reason for hiding this comment

justinchuby commented Mar 9, 2023

justinchuby commented Mar 9, 2023

justinchuby commented Mar 10, 2023

pytorchmergebot commented Mar 10, 2023

Merge started

pytorchmergebot commented Mar 10, 2023

Merge failed

justinchuby commented Mar 10, 2023

pytorchmergebot commented Mar 10, 2023

Merge started

linux-foundation-easycla bot commented Jan 12, 2023 •

edited

pytorch-bot bot commented Jan 12, 2023 •

edited

urinieto Jan 17, 2023 •

edited

urinieto commented Mar 1, 2023 •

edited