Performance issue with tf.signal.stft

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow.js): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
• TensorFlow.js installed from (npm or script link): npm i @tensorflow/tf-node
• TensorFlow.js version (use command below): 2.8.5, but master branch is affected as well
• Browser version: N/A
• Tensorflow.js Converter Version: N/A

Describe the current behavior

The tf.signal.stft function takes a 1D tensor, and calculates the FFT over a sliding window. The expected result is a 2D Tensor where all the strides contain the FFT according to the input parameters. I believe there’s a performance bottleneck in the stft.ts implementation, in the stft_ function: the input (and processed) tensor is sliced into 1D tensors, and rfft is called in a loop. The output is aggregated into an output buffer, which is concatenated and then returned.

Describe the expected behavior

The called tf.spectra.rfft function is batch-aware. There’s no reason to slice the processed tensor in stft_, you can pass windowedSignal to the rfft function, still get the same result. The benefit is, that there’s no extra memory and processing required to slice/concat the results. Also, in rfft (where fft is called, which is a kernel function) gets a boost (may the backend be in node-cpu / -gpu / webgl / etc) from the calls working on larger set of data. My measurements show at least 3-fold performance increase (see example) using tf-node-cpu.

Please consider reviewing the sfft_ function.

Standalone code to reproduce the issue

https://gist.github.com/harangp/b4c21a6c1ce9e2c8b0c1da804dde784d

Other info / logs

Output on my machine (i5 7th gen running at 3Ghz):

Standard aproach ----------------
2021-01-31T18:17:31.568Z 'ellapsed time:' 110
2021-01-31T18:17:31.650Z 'ellapsed time:' 80
2021-01-31T18:17:31.790Z 'ellapsed time:' 139
2021-01-31T18:17:31.847Z 'ellapsed time:' 56
2021-01-31T18:17:31.903Z 'ellapsed time:' 55
2021-01-31T18:17:31.961Z 'ellapsed time:' 57
2021-01-31T18:17:32.010Z 'ellapsed time:' 48
2021-01-31T18:17:32.059Z 'ellapsed time:' 48
2021-01-31T18:17:32.113Z 'ellapsed time:' 53
2021-01-31T18:17:32.161Z 'ellapsed time:' 47
2021-01-31T18:17:32.210Z 'ellapsed time:' 48
2021-01-31T18:17:32.262Z 'ellapsed time:' 51
2021-01-31T18:17:32.314Z 'ellapsed time:' 51
2021-01-31T18:17:32.362Z 'ellapsed time:' 47
2021-01-31T18:17:32.415Z 'ellapsed time:' 52
2021-01-31T18:17:32.465Z 'ellapsed time:' 49
2021-01-31T18:17:32.510Z 'ellapsed time:' 44
2021-01-31T18:17:32.552Z 'ellapsed time:' 42
2021-01-31T18:17:32.607Z 'ellapsed time:' 54
2021-01-31T18:17:32.655Z 'ellapsed time:' 48
Sped up approach ----------------
2021-01-31T18:17:32.667Z 'ellapsed time:' 11
2021-01-31T18:17:32.685Z 'ellapsed time:' 9
2021-01-31T18:17:32.696Z 'ellapsed time:' 10
2021-01-31T18:17:32.717Z 'ellapsed time:' 9
2021-01-31T18:17:32.728Z 'ellapsed time:' 10
2021-01-31T18:17:32.746Z 'ellapsed time:' 12
2021-01-31T18:17:32.768Z 'ellapsed time:' 9
2021-01-31T18:17:32.779Z 'ellapsed time:' 11
2021-01-31T18:17:32.799Z 'ellapsed time:' 11
2021-01-31T18:17:32.808Z 'ellapsed time:' 9
2021-01-31T18:17:32.828Z 'ellapsed time:' 11
2021-01-31T18:17:32.848Z 'ellapsed time:' 12
2021-01-31T18:17:32.871Z 'ellapsed time:' 18
2021-01-31T18:17:32.884Z 'ellapsed time:' 12
2021-01-31T18:17:32.903Z 'ellapsed time:' 9
2021-01-31T18:17:32.915Z 'ellapsed time:' 12
2021-01-31T18:17:32.938Z 'ellapsed time:' 11
2021-01-31T18:17:32.964Z 'ellapsed time:' 11
2021-01-31T18:17:32.984Z 'ellapsed time:' 12
2021-01-31T18:17:32.993Z 'ellapsed time:' 9
Does it return the same results?
Tensor
    true

[harangp]

@harangp @JPery Awesome work for great perf improvement, any of you want to contribute the change to the library? thanks!

Yes, happy to (will be first time, but one must start somewhere).

[google-ml-butler[bot]]

Are you satisfied with the resolution of your issue? Yes No