|
|
import torch
|
|
|
from scipy.signal import get_window
|
|
|
|
|
|
from torch import nn
|
|
|
|
|
|
'''
|
|
|
class LambdaOverlapAdd(torch.nn.Module):
|
|
|
"""Overlap-add with lambda transform on segments.
|
|
|
|
|
|
Segment input signal, apply lambda function (a neural network for example)
|
|
|
and combine with OLA.
|
|
|
|
|
|
Args:
|
|
|
nnet (callable): Function to apply to each segment.
|
|
|
n_src (int): Number of sources in the output of nnet.
|
|
|
window_size (int): Size of segmenting window.
|
|
|
hop_size (int): Segmentation hop size.
|
|
|
window (str): Name of the window (see scipy.signal.get_window) used
|
|
|
for the synthesis.
|
|
|
reorder_chunks (bool): Whether to reorder each consecutive segment.
|
|
|
This might be useful when `nnet` is permutation invariant, as
|
|
|
source assignements might change output channel from one segment
|
|
|
to the next (in classic speech separation for example).
|
|
|
Reordering is performed based on the correlation between
|
|
|
the overlapped part of consecutive segment.
|
|
|
|
|
|
Examples:
|
|
|
>>> from asteroid_test import ConvTasNet
|
|
|
>>> nnet = ConvTasNet(n_src=2)
|
|
|
>>> continuous_nnet = LambdaOverlapAdd(
|
|
|
>>> nnet=nnet,
|
|
|
>>> n_src=2,
|
|
|
>>> window_size=64000,
|
|
|
>>> hop_size=None,
|
|
|
>>> window="hanning",
|
|
|
>>> reorder_chunks=True,
|
|
|
>>> enable_grad=False,
|
|
|
>>> )
|
|
|
>>> wav = torch.randn(1, 1, 500000)
|
|
|
>>> out_wavs = continuous_nnet.forward(wav)
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
nnet,
|
|
|
n_src,
|
|
|
window_size,
|
|
|
hop_size=None,
|
|
|
window="hanning",
|
|
|
reorder_chunks=True,
|
|
|
enable_grad=False,
|
|
|
):
|
|
|
super().__init__()
|
|
|
assert window_size % 2 == 0, "Window size must be even"
|
|
|
|
|
|
self.nnet = nnet
|
|
|
self.window_size = window_size
|
|
|
self.hop_size = hop_size if hop_size is not None else window_size // 2
|
|
|
self.n_src = n_src
|
|
|
|
|
|
if window:
|
|
|
window = get_window(window, self.window_size).astype("float32")
|
|
|
window = torch.from_numpy(window)
|
|
|
self.use_window = True
|
|
|
else:
|
|
|
self.use_window = False
|
|
|
|
|
|
self.register_buffer("window", window)
|
|
|
self.reorder_chunks = reorder_chunks
|
|
|
self.enable_grad = enable_grad
|
|
|
|
|
|
def ola_forward(self, x):
|
|
|
"""Heart of the class: segment signal, apply func, combine with OLA."""
|
|
|
|
|
|
assert x.ndim == 3
|
|
|
|
|
|
batch, channels, n_frames = x.size()
|
|
|
# Overlap and add:
|
|
|
# [batch, chans, n_frames] -> [batch, chans, win_size, n_chunks]
|
|
|
unfolded = torch.nn.functional.unfold(
|
|
|
x.unsqueeze(-1),
|
|
|
kernel_size=(self.window_size, 1),
|
|
|
padding=(self.window_size, 0),
|
|
|
stride=(self.hop_size, 1),
|
|
|
)
|
|
|
|
|
|
out = []
|
|
|
n_chunks = unfolded.shape[-1]
|
|
|
for frame_idx in range(n_chunks): # for loop to spare memory
|
|
|
frame = self.nnet(unfolded[..., frame_idx])
|
|
|
# user must handle multichannel by reshaping to batch
|
|
|
if frame_idx == 0:
|
|
|
assert frame.ndim == 3, "nnet should return (batch, n_src, time)"
|
|
|
assert frame.shape[1] == self.n_src, "nnet should return (batch, n_src, time)"
|
|
|
frame = frame.reshape(batch * self.n_src, -1)
|
|
|
|
|
|
if frame_idx != 0 and self.reorder_chunks:
|
|
|
# we determine best perm based on xcorr with previous sources
|
|
|
frame = _reorder_sources(
|
|
|
frame, out[-1], self.n_src, self.window_size, self.hop_size
|
|
|
)
|
|
|
|
|
|
if self.use_window:
|
|
|
frame = frame * self.window
|
|
|
else:
|
|
|
frame = frame / (self.window_size / self.hop_size)
|
|
|
out.append(frame)
|
|
|
|
|
|
out = torch.stack(out).reshape(n_chunks, batch * self.n_src, self.window_size)
|
|
|
out = out.permute(1, 2, 0)
|
|
|
|
|
|
out = torch.nn.functional.fold(
|
|
|
out,
|
|
|
(n_frames, 1),
|
|
|
kernel_size=(self.window_size, 1),
|
|
|
padding=(self.window_size, 0),
|
|
|
stride=(self.hop_size, 1),
|
|
|
)
|
|
|
return out.squeeze(-1).reshape(batch, self.n_src, -1)
|
|
|
|
|
|
def forward(self, x):
|
|
|
"""Forward module: segment signal, apply func, combine with OLA.
|
|
|
|
|
|
Args:
|
|
|
x (:class:`torch.Tensor`): waveform signal of shape (batch, 1, time).
|
|
|
|
|
|
Returns:
|
|
|
:class:`torch.Tensor`: The output of the lambda OLA.
|
|
|
"""
|
|
|
# Here we can do the reshaping
|
|
|
with torch.autograd.set_grad_enabled(self.enable_grad):
|
|
|
olad = self.ola_forward(x)
|
|
|
return olad
|
|
|
|
|
|
|
|
|
def _reorder_sources(
|
|
|
current: torch.FloatTensor,
|
|
|
previous: torch.FloatTensor,
|
|
|
n_src: int,
|
|
|
window_size: int,
|
|
|
hop_size: int,
|
|
|
):
|
|
|
"""
|
|
|
Reorder sources in current chunk to maximize correlation with previous chunk.
|
|
|
Used for Continuous Source Separation. Standard dsp correlation is used
|
|
|
for reordering.
|
|
|
|
|
|
|
|
|
Args:
|
|
|
current (:class:`torch.Tensor`): current chunk, tensor
|
|
|
of shape (batch, n_src, window_size)
|
|
|
previous (:class:`torch.Tensor`): previous chunk, tensor
|
|
|
of shape (batch, n_src, window_size)
|
|
|
n_src (:class:`int`): number of sources.
|
|
|
window_size (:class:`int`): window_size, equal to last dimension of
|
|
|
both current and previous.
|
|
|
hop_size (:class:`int`): hop_size between current and previous tensors.
|
|
|
|
|
|
Returns:
|
|
|
current:
|
|
|
|
|
|
"""
|
|
|
batch, frames = current.size()
|
|
|
current = current.reshape(-1, n_src, frames)
|
|
|
previous = previous.reshape(-1, n_src, frames)
|
|
|
|
|
|
overlap_f = window_size - hop_size
|
|
|
|
|
|
def reorder_func(x, y):
|
|
|
x = x[..., :overlap_f]
|
|
|
y = y[..., -overlap_f:]
|
|
|
# Mean normalization
|
|
|
x = x - x.mean(-1, keepdim=True)
|
|
|
y = y - y.mean(-1, keepdim=True)
|
|
|
# Negative mean Correlation
|
|
|
return -torch.sum(x.unsqueeze(1) * y.unsqueeze(2), dim=-1)
|
|
|
|
|
|
# We maximize correlation-like between previous and current.
|
|
|
pit = PITLossWrapper(reorder_func)
|
|
|
current = pit(current, previous, return_est=True)[1]
|
|
|
return current.reshape(batch, frames)
|
|
|
'''
|
|
|
|
|
|
|
|
|
class DualPathProcessing(nn.Module):
|
|
|
"""Perform Dual-Path processing via overlap-add as in DPRNN [1].
|
|
|
|
|
|
Args:
|
|
|
chunk_size (int): Size of segmenting window.
|
|
|
hop_size (int): segmentation hop size.
|
|
|
|
|
|
References:
|
|
|
[1] "Dual-path RNN: efficient long sequence modeling for
|
|
|
time-domain single-channel speech separation", Yi Luo, Zhuo Chen
|
|
|
and Takuya Yoshioka. https://arxiv.org/abs/1910.06379
|
|
|
"""
|
|
|
|
|
|
def __init__(self, chunk_size, hop_size):
|
|
|
super(DualPathProcessing, self).__init__()
|
|
|
self.chunk_size = chunk_size
|
|
|
self.hop_size = hop_size
|
|
|
self.n_orig_frames = None
|
|
|
|
|
|
def unfold(self, x):
|
|
|
"""Unfold the feature tensor from
|
|
|
|
|
|
(batch, channels, time) to (batch, channels, chunk_size, n_chunks).
|
|
|
|
|
|
Args:
|
|
|
x: (:class:`torch.Tensor`): feature tensor of shape (batch, channels, time).
|
|
|
|
|
|
Returns:
|
|
|
x: (:class:`torch.Tensor`): spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
|
|
|
"""
|
|
|
|
|
|
batch, chan, frames = x.size()
|
|
|
assert x.ndim == 3
|
|
|
self.n_orig_frames = x.shape[-1]
|
|
|
unfolded = torch.nn.functional.unfold(
|
|
|
x.unsqueeze(-1),
|
|
|
kernel_size=(self.chunk_size, 1),
|
|
|
padding=(self.chunk_size, 0),
|
|
|
stride=(self.hop_size, 1),
|
|
|
)
|
|
|
|
|
|
return unfolded.reshape(
|
|
|
batch, chan, self.chunk_size, -1
|
|
|
)
|
|
|
|
|
|
def fold(self, x, output_size=None):
|
|
|
"""Folds back the spliced feature tensor.
|
|
|
|
|
|
Input shape (batch, channels, chunk_size, n_chunks) to original shape
|
|
|
(batch, channels, time) using overlap-add.
|
|
|
|
|
|
Args:
|
|
|
x: (:class:`torch.Tensor`): spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
output_size: (int, optional): sequence length of original feature tensor.
|
|
|
If None, the original length cached by the previous call of `unfold`
|
|
|
will be used.
|
|
|
|
|
|
Returns:
|
|
|
x: (:class:`torch.Tensor`): feature tensor of shape (batch, channels, time).
|
|
|
|
|
|
.. note:: `fold` caches the original length of the pr
|
|
|
|
|
|
"""
|
|
|
output_size = output_size if output_size is not None else self.n_orig_frames
|
|
|
|
|
|
batch, chan, chunk_size, n_chunks = x.size()
|
|
|
to_unfold = x.reshape(batch, chan * self.chunk_size, n_chunks)
|
|
|
x = torch.nn.functional.fold(
|
|
|
to_unfold,
|
|
|
(output_size, 1),
|
|
|
kernel_size=(self.chunk_size, 1),
|
|
|
padding=(self.chunk_size, 0),
|
|
|
stride=(self.hop_size, 1),
|
|
|
)
|
|
|
|
|
|
x /= self.chunk_size / self.hop_size
|
|
|
|
|
|
return x.reshape(batch, chan, self.n_orig_frames)
|
|
|
|
|
|
@staticmethod
|
|
|
def intra_process(x, module):
|
|
|
"""Performs intra-chunk processing.
|
|
|
|
|
|
Args:
|
|
|
x (:class:`torch.Tensor`): spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
module (:class:`torch.nn.Module`): module one wish to apply to each chunk
|
|
|
of the spliced feature tensor.
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
x (:class:`torch.Tensor`): processed spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
|
|
|
.. note:: the module should have the channel first convention and accept
|
|
|
a 3D tensor of shape (batch, channels, time).
|
|
|
"""
|
|
|
|
|
|
|
|
|
batch, channels, chunk_size, n_chunks = x.size()
|
|
|
|
|
|
x = x.transpose(1, -1).reshape(batch * n_chunks, chunk_size, channels).transpose(1, -1)
|
|
|
x = module(x)
|
|
|
x = x.reshape(batch, n_chunks, channels, chunk_size).transpose(1, -1).transpose(1, 2)
|
|
|
return x
|
|
|
|
|
|
@staticmethod
|
|
|
def inter_process(x, module):
|
|
|
"""Performs inter-chunk processing.
|
|
|
|
|
|
Args:
|
|
|
x (:class:`torch.Tensor`): spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
module (:class:`torch.nn.Module`): module one wish to apply between
|
|
|
each chunk of the spliced feature tensor.
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
x (:class:`torch.Tensor`): processed spliced feature tensor of shape
|
|
|
(batch, channels, chunk_size, n_chunks).
|
|
|
|
|
|
.. note:: the module should have the channel first convention and accept
|
|
|
a 3D tensor of shape (batch, channels, time).
|
|
|
"""
|
|
|
|
|
|
batch, channels, chunk_size, n_chunks = x.size()
|
|
|
x = x.transpose(1, 2).reshape(batch * chunk_size, channels, n_chunks)
|
|
|
x = module(x)
|
|
|
x = x.reshape(batch, chunk_size, channels, n_chunks).transpose(1, 2)
|
|
|
return x
|
|
|
|
