[Frontend][PyTorch] Add: Relay stft operator #11190

ghost · 2022-04-29T18:46:06Z

This PR adds the stft, amax, amin
torch.stft

# This is run using 4 cpu cores for 100 iterations average inference time per frame
pt Elapsed time: 0.1826983118057251
graph_runtime Elapsed time: 0.11588997602462768

masahi · 2022-04-29T20:06:14Z

@jsheng-jian Thanks, looks like topi/stft.py and topi/cuda/stft.py are missing?

I'm very curious how you implemented stft (without fft in TVM)!

python/setup.py

ghost · 2022-05-01T04:09:43Z

topi/stft.py and topi/cuda/stft.py are added, could you take another look? @masahi

masahi · 2022-05-02T05:03:45Z

Please address the CI issue, there is a warning error from doc.

masahi · 2022-05-03T04:11:07Z

python/tvm/relay/op/transform.py

+    win_length : int
+        The size of window frame and STFT filter
+    window : relay.Expr
+        A 1-D tensor window frame


In PyTorch, window argument is optional. So shouldn't we support that too?

Added optional window argument.

masahi · 2022-05-03T04:11:44Z

python/tvm/relay/op/transform.py

+    Returns
+    -------
+    output : relay.Expr
+        Tensor containing the STFT result


Document the output shape. I had to read the type rel to see how the output shape looks like.

masahi · 2022-05-03T04:12:32Z

python/tvm/topi/cuda/stft.py

+
+            with ib.for_range(0, output_ptr.shape[0]) as batch:
+                with ib.for_range(0, output_ptr.shape[1]) as row:
+                    with ib.for_range(0, output_ptr.shape[2]) as col:


This looks weird. You try to parallelize over the batch dim but nothing is parallelized.

Updated Cuda implementation.

masahi · 2022-05-03T04:13:34Z

python/tvm/topi/stft.py

+
+        with ib.for_range(0, output_ptr.shape[0]) as batch:
+            # https://librosa.org/doc/0.7.2/_modules/librosa/core/spectrum.html#stft
+            with ib.for_range(0, output_ptr.shape[1], kind="parallel") as row:


fuse the outer loop to have one big parallel loop.

masahi · 2022-05-03T04:14:59Z

python/tvm/topi/stft.py

+
+    output_buf = tir.decl_buffer(output_shape, data.dtype, "output_buf")
+    loop_kind = "vectorize"
+    if hasattr(output_shape[2], "name") and output_shape[2].name == "any_dim":


if isinstance(output_shape[2], tir.expr.Any)

The type is tir.expr.SizeVar, updated.

masahi · 2022-05-03T04:15:41Z

tests/python/frontend/pytorch/test_forward.py

+    verify_trace_model(test_fn(3, 3, 3, False, "reflect", False, True), [input, window], targets)
+    window = torch.tensor([1, 3], dtype=torch.int32)
+    verify_trace_model(test_fn(2, 1, 2, False, "reflect", False, True), [input, window], targets)
+


Please add a test for window=None case.

masahi · 2022-05-03T04:15:53Z

tests/python/relay/test_op_level3.py

-            tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
-        relay.backend.te_compiler.get().clear()
-
-


Remove this diff

masahi · 2022-05-03T04:16:39Z

python/tvm/relay/op/transform.py

+def stft(data, n_fft, hop_length, win_length, window, normalized, onesided):
+    """
+    The STFT computes the Fourier transform of short overlapping windows of the input.
+    This giving frequency components of the signal as they change over time.


giving -> gives

Fix the same typo in other files too.

masahi · 2022-05-04T04:42:08Z

tests/python/relay/test_op_level3.py

+        )
+
+
+def verify_func2(target, dev, func, data, ref_res, rtol=1e-5, atol=1e-7, kinds=["vm"]):


Why do you need this? It looks identical to verify_func.

I want to expost rtol atol and kinds, I can update the original function instead of creating a new one?

Yeah that would be better

masahi · 2022-05-04T04:43:12Z

python/tvm/topi/cuda/stft.py

+    win_length,
+    window,
+    normalized,
+    onesided,  # pylint: disable=unused-argument


You can remove # pylint: disable=unused-argument and add unused-argument to L17.

masahi · 2022-05-04T04:43:22Z

python/tvm/topi/cuda/stft.py

+        win_length,
+        window_ptr,
+        normalized,
+        onesided,  # pylint: disable=unused-argument


masahi · 2022-05-04T04:43:48Z

python/tvm/topi/stft.py

+    win_length,
+    window,
+    normalized,
+    onesided,  # pylint: disable=unused-argument


same as the comment in topi/cuda/stft.py

masahi · 2022-05-04T04:44:12Z

python/tvm/topi/stft.py

+        win_length,
+        window_ptr,
+        normalized,
+        onesided,  # pylint: disable=unused-argument


ghost · 2022-05-05T18:34:41Z

The CI failure seems to be unrelated.

ghost · 2022-05-06T18:37:42Z

It looks the CI failure is resolved in main branch, do I need to rebase my changes?

masahi · 2022-05-06T22:15:42Z

Yes, please send another job

* Add: Relay stft operator * fix doc * address PR comments * address addtional comments

ghost changed the title ~~Add: Relay stft operator~~ [Frontend][PyTorch] Add: Relay stft operator Apr 29, 2022

masahi reviewed Apr 29, 2022

View reviewed changes

python/setup.py Outdated Show resolved Hide resolved

masahi requested changes May 3, 2022

View reviewed changes

masahi reviewed May 4, 2022

View reviewed changes

masahi approved these changes May 5, 2022

View reviewed changes

Jian Sheng added 4 commits May 6, 2022 17:52

Add: Relay stft operator

eca2cac

fix doc

c0f3f4b

address PR comments

a371d0d

address addtional comments

5c58e8d

masahi merged commit a3d75ae into apache:main May 7, 2022

shtinsa pushed a commit to Deelvin/tvm that referenced this pull request May 17, 2022

[Frontend][PyTorch] Add: Relay stft operator (apache#11190)

278e3e5

* Add: Relay stft operator * fix doc * address PR comments * address addtional comments

SebastianBoblest pushed a commit to SebastianBoblest/tvm that referenced this pull request May 27, 2022

[Frontend][PyTorch] Add: Relay stft operator (apache#11190)

ff1a4fb

* Add: Relay stft operator * fix doc * address PR comments * address addtional comments

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

masahi mentioned this pull request Sep 29, 2022

[PyTorch] NotImplementedError: The following operators are not implemented: ['aten::stft'] #8877

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Frontend][PyTorch] Add: Relay stft operator #11190

[Frontend][PyTorch] Add: Relay stft operator #11190

ghost commented Apr 29, 2022 •

edited by ghost

Loading

masahi commented Apr 29, 2022 •

edited

Loading

ghost commented May 1, 2022

masahi commented May 2, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 3, 2022

ghost May 4, 2022

masahi May 4, 2022

ghost May 4, 2022

masahi May 4, 2022

masahi May 4, 2022

ghost May 4, 2022

masahi May 4, 2022

ghost May 4, 2022

masahi May 4, 2022

ghost May 4, 2022

masahi May 4, 2022

ghost May 4, 2022

ghost commented May 5, 2022

ghost commented May 6, 2022

masahi commented May 6, 2022

		tvm.testing.assert_allclose(op_res.numpy(), ref_res, rtol=1e-5)
		relay.backend.te_compiler.get().clear()

		)


		def verify_func2(target, dev, func, data, ref_res, rtol=1e-5, atol=1e-7, kinds=["vm"]):

[Frontend][PyTorch] Add: Relay stft operator #11190

[Frontend][PyTorch] Add: Relay stft operator #11190

Conversation

ghost commented Apr 29, 2022 • edited by ghost Loading

masahi commented Apr 29, 2022 • edited Loading

ghost commented May 1, 2022

masahi commented May 2, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ghost commented May 5, 2022

ghost commented May 6, 2022

masahi commented May 6, 2022

ghost commented Apr 29, 2022 •

edited by ghost

Loading

masahi commented Apr 29, 2022 •

edited

Loading