Skip to content

Add stft option to align window for center = false#145324

Merged
jackzhxng merged 5 commits intojz/stft-old-fcfrom
jz/stft
Jan 30, 2025
Merged

Add stft option to align window for center = false#145324
jackzhxng merged 5 commits intojz/stft-old-fcfrom
jz/stft

Conversation

@jackzhxng
Copy link
Contributor

@jackzhxng jackzhxng commented Jan 21, 2025

Adds a flag for aligning the start of the window to the start of the signal when center = false (aka window-based padding). The same flag was proposed a while ago for Librosa as well.

For internal reasons, we need to add this behavior to the op and this flag allows us to do so while preserving bc compatibility

Pr chain:

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Jan 21, 2025
@pytorch-bot
Copy link

pytorch-bot bot commented Jan 21, 2025

🔗 Helpful Links

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

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

✅ No Failures

As of commit 80e461e with merge base b2c89bc (image):
💚 Looks good so far! There are no failures yet. 💚

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

@jackzhxng jackzhxng removed the release notes: onnx torch.onnx related changes that should show up in the release notes label Jan 21, 2025
@github-actions
Copy link
Contributor

Attention! native_functions.yaml was changed

If you are adding a new function or defaulted argument to native_functions.yaml, you cannot use it from pre-existing Python frontend code until our FC window passes (two weeks). Split your PR into two PRs, one which adds the new C++ functionality, and one that makes use of it from Python, and land them two weeks apart. See https://github.com/pytorch/pytorch/wiki/PyTorch's-Python-Frontend-Backward-and-Forward-Compatibility-Policy#forwards-compatibility-fc for more info.


Caused by:

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Jan 21, 2025
@jackzhxng jackzhxng force-pushed the jz/stft branch 4 times, most recently from 254fdc5 to e386693 Compare January 23, 2025 00:47
@jackzhxng jackzhxng changed the base branch from main to jz/stft-old-fc January 23, 2025 00:50
@jackzhxng jackzhxng force-pushed the jz/stft branch 2 times, most recently from e264167 to 4536cfd Compare January 23, 2025 04:46
@jackzhxng jackzhxng force-pushed the jz/stft-old-fc branch 2 times, most recently from e9c4dff to aa60e40 Compare January 23, 2025 04:51
@ezyang
Copy link
Contributor

ezyang commented Jan 25, 2025

cc @philgzl, I would highly value your opinion here!

@philgzl
Copy link

philgzl commented Jan 25, 2025

I think it looks fine. As discussed here it could be that padding after/before striding is equivalent after this fix. Although in the code snippet I wrote there the results are not equal (but again I didn't investigate much).


# Overload without center & pad mode, needed for forward-compatibility
- func: stft(Tensor self, int n_fft, int? hop_length=None, int? win_length=None, Tensor? window=None, bool normalized=False, bool? onesided=None, bool? return_complex=None) -> Tensor
- func: stft(Tensor self, int n_fft, int? hop_length=None, int? win_length=None, Tensor? window=None, bool normalized=False, bool? onesided=None, bool? return_complex=None, bool? align_to_window=None) -> Tensor
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked in with @rgommers and his suggestion for another way to specify the mode is to instead do str? align_mode=None which will allow for different alignment modes in the future. This would then match the API proposal in librosa: librosa/librosa#596 (comment)

Ralf also said he'd ping a scipy expert to come take a look :)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that I'm not super attached to my suggestion, but in general when one has multiple modes a string option is cleaner than multiple boolean options that interact.

Ralf also said he'd ping a scipy expert to come take a look :)

@DietBru I thought this may interest you. I looked at scipy.signal.ShortTimeFFT as a reference, and don't think we have this mode there - but I could be wrong, it's not trivial to be sure given that there are quite a few modes and they don't match between SciPy and PyTorch.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taking the liberty to chip in here.

don't think we have this mode there

Probably because align_to_window=True should be the default behavior. Without it, the first samples are discarded when center=False and n_fft > win_length. I don't see why anyone wouldn't use this other than preserving BC.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In that case, it might not be horrible to have an awful name, with the intention of flipping the default after a deprecation cycle. IDK, @dvorjackz, I'll defer to your choice.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Concerning ShortTimeFFT:
The default behavior is to center the first window slice over the first signal sample, as illustrated here. This also incorporates a phase shift, to ensure that the phase of each STFT slice is correct. Furthermore, ShortTimeFFT.stft() provides a parameter k_offset to shift the input signal.
In summary: The parameter k_offset and the phase shift allow to arbitrarily set window alignments on the signal, which makes the ShortTimeFFT much more flexible.

I am not sure I can be of much help here. The challenge with aligning the window is to ensure consistent behavior at the borders—especially when it is desired to recreate the original signal with the ISTFT (one of the main motivations for the ShortTimeFFT implementation) .

@justinchuby
Copy link
Collaborator

Just checking, is there a OpInfo for this op that needs to be updated?

@ezyang
Copy link
Contributor

ezyang commented Jan 28, 2025

Yes please

@jackzhxng
Copy link
Contributor Author

Just checking, is there a OpInfo for this op that needs to be updated?

None yet it seems

Copy link
Contributor

@iseeyuan iseeyuan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thanks @dvorjackz !

pytorchmergebot pushed a commit that referenced this pull request Jan 30, 2025
Long overdue follow-up on https://github.com/pytorch/pytorch/pull/73432/files#diff-5f3d4caa0693a716fc46fd7f6339312f1b5f0bf89e3a3ff58e9dc13a9486b17aR719

Onnx stft doesn't support centering, [and all of the existing tests are for center = False](https://github.com/pytorch/pytorch/blob/main/test/onnx/test_pytorch_onnx_onnxruntime.py#L8026). I will open a follow-up issue to address this, this is just a nice-to-have.

Pr chain:
- -> [Advance past fc window for stft center #145437](#145437)
- [Add stft option to align window for center = false #145324](#145324)
- [Add istft option to align window for center = false](#145510)
Pull Request resolved: #145437
Approved by: https://github.com/justinchuby, https://github.com/iseeyuan
@jackzhxng jackzhxng merged commit ca100b1 into jz/stft-old-fc Jan 30, 2025
101 checks passed
@jackzhxng jackzhxng deleted the jz/stft branch January 30, 2025 22:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

release notes: onnx torch.onnx related changes that should show up in the release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants