normalize STFT output by magnitude

Question

I am using torch.stft() to generate spectrograms for neural networks and come across the below code.

S = torch.stft(
        input=y, # shape(1 x num_samples)
        n_fft=self.n_fft,
        hop_length=self.hop_length,
        window=self.window,
        center=self.center,
        onesided=True,
        normalized=True
    )

And torchaudio has the below implementation:

    if normalized:
       S /= self.window.pow(2).sum().sqrt()

I vaguely know that this normalization is for energy conservation (Parseval's theorem) to restore the energy lost when applying windowing but I could not find more detailed explanation online regarding why the formular is like this.

I would also like to know if it applies to all kinds of window functions, as I also other posts showing different ways of normalizing the energy. (I assume it is generic. Otherwise, torchaudio wouldn't have used it.)

After the stft transformation, I also saw people using

S = S.pow(2).sum(-1) 
return S

The output of STFT (torch real tensor S) has the last dimension containing real and imaginary part.

Is pow(2).sum(-1) again some normalization or does it have something to do with power spectrum? (Sorry, I am a beginner in signal processing.) I don't understand what it is for. And why we don't need sqrt(s) before return s?

It would be great if you could give me some hints regarding these two operations. Many thanks in advance!

Answer 1

You are spot on with the reference to conservation of energy. Moreover, that scaling term renders an unbiased estimate for power spectral density estimation for the modified periodogram.

Following the normalized STFT calculation with

S = S.pow(2).sum(-1) 
return S

is simply the magnitude of the STFT and a summation along the last dimension (presumably time), which is Welch's method for PSD estimation scaled by the number of time bins.

And why we don't need sqrt(s) before return s?

Page 6 of the linked report lists the definition of Welch's method in terms of the discrete Fourier transform (DFT) which is where this code comes from. If it helps your intuition, recall that if a measured signal has units volts, the PSD will have units volts^2/Hz.