InverseSpectrogram output is shorter than Spectrogram input for same kwargs

[drscotthawley]

🐛 Describe the bug

Hi, I've seen this issue about when center=False, but even with center=True, I'm finding that passing a waveform into torchaudio.transforms.Spectrogram(power=None, n_fft=n_fft, hop_length=hop_length) and then immediatley running that result through torchaudio.transforms.InverseSpectrogram(n_fft=n_fft, hop_length=hop_length) results in an output waveform that is shorter than the input waveform.

I can imagine that the output might be longer due to zero-padding, but I'm not figuring out why it's shorter.

Changing to center=False did not fix the problem, rather it resulted in a runtime error.

Pre-padding the input waveform up to a power of 2 in length removes the issue, but... I would've thought this would be handled automatically.

Example:

from torchaudio import transforms as T

waveform, sr = torchaudio.load('example.wav')
print("waveform.shape = ",waveform.shape)

n_fft=1024   
hop_length=256
center=True
stft = T.Spectrogram(power=None, n_fft=n_fft, hop_length=hop_length, center=center)
istft = T.InverseSpectrogram(    n_fft=n_fft, hop_length=hop_length, center=center)

spec = stft(waveform)
recon = istft(spec)
print("spec.shape, recon.shape = ",spec.shape, recon.shape)

assert waveform.shape == recon.shape, \
    f"Expected waveform.shape ({waveform.shape}) == recon.shape ({recon.shape})"

Output:

waveform.shape = torch.Size([1, 55728])
spec.shape, recon.shape =  torch.Size([1, 513, 218]) torch.Size([1, 55552])
AssertionError: Expected waveform.shape (torch.Size([1, 55728])) == recon.shape (torch.Size([1, 55552]))

Again, I can imagine why recon.shape[-1] might be greater than 55728, but I'm having a hard time understanding why it'd be less.

Can you help me?

Also, note that current documenation for InverseSpectrogram doesn't say anything regarding any length kwarg that was mentioned in the closed Issue I referenced earlier.

Thanks for your work on this!

Versions

$ python collect_env.py 
Collecting environment information...
PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.24.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-1023-aws-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
[NOTE: censoring output manually:  removing additional GPU's from list]

Nvidia driver version: 510.73.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] nwt-pytorch==0.0.4
[pip3] perceiver-pytorch==0.8.7
[pip3] pytorch-lightning==1.8.6
[pip3] torch==1.13.1
[pip3] torchaudio==0.13.1
[pip3] torchmetrics==0.11.0
[pip3] torchvision==0.14.1
[pip3] vector-quantize-pytorch==0.10.14
[conda] Could not collect

[drscotthawley]

One more note: my intended use case was to make a simple "autoencoder" using a Spectrogram and an InverseSpectrogram. Perhaps I was just calling something wrong?

As I noted above, padding to a power of two beforehand can serve as a workaround, but... was not "expected behavior".

Here I'll share my "autoencoder" with this workaround in place.

def next_power_of_2(x:int):  
    return 1 if x == 0 else 2**(x - 1).bit_length()

def zero_pad_po2(x):
    new_shape = list(x.shape)
    new_shape[-1] = next_power_of_2(new_shape[-1])
    new_x = torch.zeros(new_shape)
    new_x[:,:x.shape[-1]] = x
    return new_x

class SpectrogramAE(nn.Module):
    "'AutoEncoder' using (complex) spectrogram. See torchaudio.Spectrogram & InverseSpectrogram for kwarg info"
    def __init__(self,
        n_fft=1024,   
        hop_length=256,
        center=True,
        zero_pad=True,
    ):
        super().__init__()
        self.stft = T.Spectrogram(power=None, n_fft=n_fft, hop_length=hop_length, center=center)
        self.istft = T.InverseSpectrogram(    n_fft=n_fft, hop_length=hop_length, center=center)
        self.orig_len, self.zero_pad = None, zero_pad 
        
    def encode(self, waveform: torch.Tensor) -> torch.Tensor:
        self.orig_len = waveform.shape[-1]
        return self.stft(zero_pad_po2(waveform)) if self.zero_pad else self.stft(waveform)
        
    def decode(self, reps: torch.Tensor) -> torch.Tensor:
        recon = self.istft(reps)
        if self.orig_len is not None: recon = recon[:,:self.orig_len]
        return recon
    
    def forward(self, waveform: torch.Tensor):
        spec = self.encode(waveform)
        recon = self.decode(spec)
        return spec, recon

[nateanl]

Hi @drscotthawley, thanks for sharing the issue. I can reproduce it and know why the length is shorter.
Basically when calling Spectrogram with center=True, the function will pad n_fft//2 to both the beginning and end of the waveform. However, since torch.stft only uses number of samples in power of 2 and discards the remaining samples, not all padded values are used in the spectrogram.
Then when calling InverseSpectrogram with center=True and with length argument as None, torch.istft assumes the inversed waveform needs to trim n_fft//2 samples at both the beginning and the end. As we have lost some samples in torch.stft, the final waveform is shorter than before.

A good solution is using length argument when calling the forward method of InverseSpectrogram. The documentation for length argument is not listed in the initialization method, but in the forward method. Do you think adding a note regarding how to use length argument in the document can help make users better understand the usage?