Torchaudio Complex Tensor Support and Migration

Overview

torchaudio has been expressing complex numbers by having an extra dimension for real-part and imaginary-part. (We will refer this format as "pseudo complex type")

waveform = torch.randn(...)
spectrogram = torchaudio.functional.spectrogram(waveform, ..., power=None)

# the last dimension represents real and imaginary parts of complex tensors.
print(spectrogram.dtype, spectrogram.shape)
>>> torch.float32, torch.Size([... ,2])

PyTorch 1.6 introduced complex Tensor type, such as torch.complex64 (torch.cfloat) and torch.complex128 (torch.cdouble). (Will be refered as "native complex type")

The natitve complex type comes with handy methods for complex operation such as abs, angle and magphase. (Please refer to the official documentation for the detail.)

Over the few coming releases, we plan to migrate torchaudio's functions and transforms to the native complex type. This issues describes the planned approaches/works/changes/timeline. If you have a question, a concern or a suggestion. Feel free to leave a comment.

Migration Stages

We will perform the migration in multiple stages. At this moment, the completion of the later migration stages re not tied with specific releases yet.

✅ Stage 0 (~ 0.8)

Up to release 0.8, torchaudio exclusively used pseudo complex type. In PyTorch 1.7, PyTorch started the adaptation of native complex type and the migration of torch.fft namespace. Because of this, torchaudio already uses native complex type in some implementations (F.vad, T.Vad, kaldi.spectrogram and kaldi.fbank) but all the user facing APIs use use pseudo complex type.

✅ Stage 1 Add support for native complex type and deprecate pseudo complex type)

Completed: PyTorch 1.9 / torchaudio 0.9

Library code change

In this stage, torchaudio will support both pseudo complex type and native complex type. This means that

• Functions that accept complex input should be able to handle both pseudo and native complex types.
• Functions that return complex values can return both pseudo and native complex types.
  • For real-to-complex functions, a new argument return_complex will be added so that users can switch the behavior.

Test code update

In addition to the above library code changes, we are going to add a set of tests to make sure that native complex types work in common use cases. This includes;

• Gradient check
• JIT
• Performance
• Distributed training
• nn.Module compatibility

✅Stage 2 (Switch to native complex type by default)

Completed: main branch. To be released as part of PyTorch 1.10 / torchaudio 0.10

The default value for return_complex is changed to True.

👉 Stage 3 (Remove the support for pseudo complex type)

In this stage, we will remove the support for pseudo complex type.

• Functions that work with complex types should handle native complex types exclusively.
• Passing pseudo complex type results in an error.
• return_complex argument added in Stage. 1 is deprecated and eventually removed.

Affected Functions

The following figure illustrates the functions that handle complex values and their dependencies.

Utility functions

F.angle, F.complex_norm, T.ComplexNorm, F.magphase

These functions are deprecated in Stage.1 and will be removed in Stage.3.
For F.angle, native complex tensors provide the angle() function.
For F.complex_norm / T.ComplexNorm, the equivalent computation can be performed with abs().pow(n).
F.magphase is a convenient function to call F.angle and F.complex_norm, therefore, this function is deprecated as well.

Real to real functions

F.griffinlim, T.GriffinLim

Changes to these functions are kept internal, therefore we can simply change the internals without disturbing the downstream users.

Complex to complex functions

F.phase_vocoder, T.TimeStretch

When adding support for native complex type, we can simplify the interface change as follow

• If the input is pseudo complex type, return pseudo complex type
• If the input is native complex type, return native complex type

Real to complex functions

F.spectrogram, T.Spectrogram

These functions return either real valued Tensor (power, energy) or complex valued Tensor (frequency representation), which depends on what power argument was provided. When power is not provided, these functions return a complex-valued Tensor. In this case, users have the option to receive the result in pseudo complex type or native complex type. return_complex argument will be added for this choice. If return_complex is True, then native complex type is returned. See #1009 for the discussion.

Timeline

Migration steps

F.angle, F.complex_norm, F.magphase and T.ComplexNorm

spectrogram = ...  # Tensor with pseudo complex type (shape == (..., 2))
angle = F.angle(spectrogram)
magnitude = F.complex_norm(spectrogram, norm=1)
power = F.complex_norm(spectrogram, norm=2)
norm = F.complex_norm(spectrogram, norm=norm)
magnitude, phase = F.magphase(spectrogram, n)

spectrogram = ...  # Tensor with pseudo complex type (shape == (..., 2))
spectrogram = torch.view_as_complex(spectrogram)  
angle = spectrogram.angle()
magnitude = spectrogram.abs()
power = spectrogram.abs().pow(2)
norm = spectrogram.abs().pow(norm)
magnitude, phase = spectrogram.abs().pow(n), spectrogram.angle()

F.phase_vocoder, T.TimeStretch

spec = ... # pseudo complex (..., 2) 

## If using functional form
spec = F.phase_vocoder(spec, ...)
## else using transform
transform = T.TimeStretch(...)
spec = transform(spec)

# Convert to native complex
spec = ... # pseudo complex (..., 2) 
# convert to native complex type
spec = torch.view_as_complex(spec)

# Perform the operation
## If using functional form
spec = F.phase_vocoder(spec, ...)
## else using transform
transform = T.TimeStretch(...)
spec = transform(spec)

# Convert back to pseudo complex type
# (If your downstream code still expects pseudo complex type)
spec = torch.view_as_real(spec)

F.spectrogram, T.Spectrogram

spec = F.spectrogram(waveform, ..., power=None)

transform = T.Spectrogram(..., power=None)
spec = transform(waveform)  # pseudo complex (..., 2) 

spec = F.spectrogram(waveform, ..., power=None, return_complex=True)

transform = T.Spectrogram(..., power=None, return_complex=True)
spec = transform(waveform)  # native complex
# If your downstream code still expects pseudo complex type
spec = torch.view_as_real(spec)

PRs - TODO (@mthrok)

Migration

Phase 1

Code Change

• F.griffinlim
  Adopt native complex dtype in griffnlim #1368
• F.phase_vocoder, T.TimeStretch
  Make F.phase_vocoder and T.TimeStretch handle complex dtype #1410
  Use complex tensors in phase_vocoder #758
• F.spectrogram, T.Spectrogram, T.MelSpectrogram
  Add return_complex to F.spectrogram and T.Spectrogram #1366
  [D] Update spectrogram to use complex #1009

Add deprecation Warnings

• F.angle, F.complex_norm, T.ComplexNorm, F.phase_vocoder, T.TimeStretch, F.spectrogram, T.Spectrogram
  Add deprecation warnings to functions for complex #1445
  Add warning about complex tensor to spectrogram #1431
• F.magphase Add deprecation warnings to magphase and ComplexNorm #1492

Phase 2

Change the default value of return_complex to True.

• F.spectrogram, T.Spectrogram [BC-Breaking] Default to native complex type when returning raw spect… #1549

Update the deprecation warnings to indicate the version of removal.

• F.angle, F.complex_norm, F.magphase, F.phase_vocoder, T.TimeStretch, T.ComplexNorm Set removal version of pseudo complex support #1553

Phase 3

Remove the support for pseudo complex type.

• Remove F.magphase [BC-Breaking] Remove deprecated F.magphase #1934
• Remove F.angle [BC-Breaking] Remove deprecated F.angle #1935
• Remove F.complex_norm and T.ComplexNorm [BC-Breaking] Remove F.complex_norm and T.ComplexNorm #1942
• Remove support for pseudo complex type from F.phase_vocoder and T.TimeStretch [BC-Breaking] Drop pseudo complex support from phase_vocoder / TimeStretch #1957
• Remove support for pseudo complex type from F.spectrogram and T.Spectrogram [BC-Breaking] Drop pseudo complex support from spectrogram #1958
  • Deprecate return_complex argument from F.spectrogram and T.Spectrogram

Surrounding works

Conjugate input tests

• After Conjugate View pytorch#54987, check if things work when the input complex Tensor is conjugate.

Autograd tests

• Add autograd test for T.GriffinLim #1421 T.GriffinLim
  • [Optional] F.griffinlim
• Add autograd test to T.TimeStretch (and F.phase_vocoder) #1420 T.TimeStretch, F.phase_vocoder
  NOTE They are not differentiable around zero.
• Add autograd test for T.Spectrogram/T.MelSpectrogram #1340 T.Spectrogram, T.MelSpectrogram
  • Investigate why T.Spectrogram is not deterministic on its backward pass. [error -> issue pytorch/#54093]
  • [Optional] F.spectrogram

Ensuring TorchScript support

Check if all the functionals/transforms are covered by TorchScript consistency test and add if missing

• F.spectrogram
• F.griffinlim
• F.phase_vocoder
• T.Spectrogram
• T.GriffinLim
• T.MelSpectrogram
• T.TimeStretch

• Save/load TorchScript object in test #1446 Add step to dump TorchScripted object in TorchScript test.

Benchmark

• F.phase_vocoder Make F.phase_vocoder and T.TimeStretch handle complex dtype #1410 x6 speed up on CPU. Not much different on GPU

cc @anjali411 @vincentqb