elixir-nx / nx_signal Goto Github PK

Stargazers

Forkers

nbw roger120981 santiago-imelio

import time
import torch


class STFT(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fft_length = 400
        self.hop_length = 160

        self.window = torch.nn.Parameter(
            torch.hann_window(self.fft_length, periodic=True), requires_grad=False
        )

    def forward(self, sample):
        stft = torch.stft(
            sample,
            self.fft_length,
            self.hop_length,
            window=self.window,
            return_complex=True,
        )

        return stft


if __name__ == "__main__":
    s = STFT().eval()
    i = torch.zeros([480000])

    start = time.time()
    result = s(i)
    print(f"took {(time.time() - start)*1000000}us")

defmodule StftTest do
  import Nx.Defn

  defn test(sample) do
    fft_length = 400
    sample_rate = 16000
    hop_length = 160

    window = NxSignal.Windows.hann(n: fft_length, is_periodic: true)

    {stft, _, _} =
      NxSignal.stft(sample, window,
        sampling_rate: sample_rate,
        fft_length: fft_length,
        overlap_length: fft_length - hop_length,
        window_padding: :reflect
      )

    stft
  end

  def test_jit() do
    sample = Nx.broadcast(0.0, {480_000})
    {time, _} = :timer.tc(fn -> Nx.Defn.jit(&test/1, compiler: EXLA).(sample) end)
    IO.puts("STFT took #{time}us")
  end
end