def __init__(self, in_channels, out_channels, skip_channels, kernel_size=3, stride=2, dilation=1, separable=False, causal=True, nonlinear='gated', norm=True, conditioning=None, enc_dim=None, enc_kernel_size=None, enc_stride=None, eps=EPS):
super().__init__()
self.kernel_size, self.stride, self.dilation = kernel_size, stride, dilation
self.separable = separable
self.nonlinear, self.norm = nonlinear, norm
# TODO: implement nonlinear & norm
if nonlinear == 'gated':
self.conv1d = GatedConv1d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, separable=separable, causal=causal, conditioning=conditioning, enc_dim=enc_dim, enc_kernel_size=enc_kernel_size, enc_stride=enc_stride)
elif nonlinear == 'relu':
raise NotImplementedError("No implementation of nonlinear ReLU.")
net = []
if separable:
net.append(SeparableConv1d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, bias=False))
else:
net.append(nn.Conv1d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, bias=False))
net.append(nn.ReLU())
self.conv1d = nn.Sequential(*net)
else:
raise ValueError("Not support {}".format(nonlinear))
if norm:
if causal:
self.norm1d = choose_layer_norm(out_channels, causal, eps=eps)
else:
self.norm1d = nn.BatchNorm1d(out_channels, eps=eps)
self.bottleneck_conv1d_output = nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=1)
self.bottleneck_conv1d_skip = nn.Conv1d(in_channels, skip_channels, kernel_size=1, stride=1)
def __init__(self,
num_features,
hidden_channels,
causal=False,
norm=True,
eps=EPS):
super().__init__()
if causal:
bidirectional = False
num_directions = 1 # uni-direction
else:
bidirectional = True
num_directions = 2 # bi-direction
self.norm = norm
self.rnn = nn.LSTM(num_features,
hidden_channels,
batch_first=False,
bidirectional=bidirectional)
self.relu = nn.ReLU()
self.fc = nn.Linear(num_directions * hidden_channels, num_features)
if self.norm:
self.norm1d = choose_layer_norm(num_features,
causal=causal,
eps=eps)
def __init__(self,
num_features,
hidden_channels=256,
skip_channels=256,
kernel_size=3,
stride=2,
dilation=1,
separable=False,
causal=True,
nonlinear=None,
norm=True,
dual_head=True,
eps=EPS):
super().__init__()
self.kernel_size, self.stride, self.dilation = kernel_size, stride, dilation
self.separable, self.causal = separable, causal
self.norm = norm
self.dual_head = dual_head
self.bottleneck_conv1d = nn.Conv1d(num_features,
hidden_channels,
kernel_size=1,
stride=1)
if nonlinear is not None:
if nonlinear == 'prelu':
self.nonlinear1d = nn.PReLU()
else:
raise ValueError("Not support {}".format(nonlinear))
self.nonlinear = True
else:
self.nonlinear = False
if norm:
self.norm1d = choose_layer_norm(hidden_channels, causal, eps=eps)
if separable:
self.separable_conv1d = DepthwiseSeparableConv1d(
hidden_channels,
num_features,
skip_channels=skip_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
causal=causal,
nonlinear=nonlinear,
norm=norm,
dual_head=dual_head,
eps=eps)
else:
if dual_head:
self.output_conv1d = nn.Conv1d(hidden_channels,
num_features,
kernel_size=kernel_size,
dilation=dilation)
self.skip_conv1d = nn.Conv1d(hidden_channels,
skip_channels,
kernel_size=kernel_size,
dilation=dilation)
def __init__(self, embed_dim, num_heads, causal=False, norm=True, eps=EPS):
super().__init__()
self.norm = norm
self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
if self.norm:
self.norm1d = choose_layer_norm(embed_dim, causal=causal, eps=eps)
def __init__(self, num_features, hidden_channels, norm=True, eps=EPS):
super().__init__()
self.num_features, self.hidden_channels = num_features, hidden_channels
num_directions = 2 # bi-direction
self.norm = norm
self.rnn = nn.LSTM(num_features, hidden_channels, batch_first=True, bidirectional=True)
self.fc = nn.Linear(num_directions*hidden_channels, num_features)
if self.norm:
self.norm1d = choose_layer_norm(num_features, causal=False, eps=eps)
def __init__(self,
num_features,
hidden_channels,
causal,
norm=True,
eps=EPS):
super().__init__()
self.norm = norm
if self.norm:
self.norm1d = choose_layer_norm(embed_dim, causal=causal, eps=eps)
self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
def __init__(self,
num_features,
bottleneck_channels=64,
hidden_channels=128,
chunk_size=100,
hop_size=50,
num_blocks=6,
causal=True,
norm=True,
mask_nonlinear='sigmoid',
n_sources=2,
eps=EPS):
super().__init__()
self.num_features, self.n_sources = num_features, n_sources
self.chunk_size, self.hop_size = chunk_size, hop_size
self.norm = norm
if self.norm:
self.norm1d = choose_layer_norm(num_features,
causal=causal,
eps=eps)
self.bottleneck_conv1d = nn.Conv1d(num_features,
bottleneck_channels,
kernel_size=1,
stride=1)
self.segment1d = Segment1d(chunk_size, hop_size)
self.dprnn = DPRNN(bottleneck_channels,
hidden_channels,
num_blocks=num_blocks,
causal=causal,
norm=norm)
self.overlap_add1d = OverlapAdd1d(chunk_size, hop_size)
self.prelu = nn.PReLU()
self.mask_conv1d = nn.Conv1d(bottleneck_channels,
n_sources * num_features,
kernel_size=1,
stride=1)
if mask_nonlinear == 'sigmoid':
self.mask_nonlinear = nn.Sigmoid()
elif mask_nonlinear == 'softmax':
self.mask_nonlinear = nn.Softmax(dim=1)
else:
raise ValueError("Cannot support {}".format(mask_nonlinear))
def __init__(self,
num_features,
bottleneck_channels=128,
hidden_channels=256,
skip_channels=128,
kernel_size=3,
num_blocks=3,
num_layers=8,
dilated=True,
separable=True,
causal=True,
nonlinear='prelu',
norm=True,
mask_nonlinear='sigmoid',
n_sources=2,
eps=EPS):
super().__init__()
self.num_features, self.n_sources = num_features, n_sources
self.norm1d = choose_layer_norm(num_features, causal=causal, eps=eps)
self.bottleneck_conv1d = nn.Conv1d(num_features,
bottleneck_channels,
kernel_size=1,
stride=1)
self.tcn = TemporalConvNet(bottleneck_channels,
hidden_channels=hidden_channels,
skip_channels=skip_channels,
kernel_size=kernel_size,
num_blocks=num_blocks,
num_layers=num_layers,
dilated=dilated,
separable=separable,
causal=causal,
nonlinear=nonlinear,
norm=norm)
self.prelu = nn.PReLU()
self.mask_conv1d = nn.Conv1d(skip_channels,
n_sources * num_features,
kernel_size=1,
stride=1)
if mask_nonlinear == 'sigmoid':
self.mask_nonlinear = nn.Sigmoid()
elif mask_nonlinear == 'softmax':
self.mask_nonlinear = nn.Softmax(dim=1)
else:
raise ValueError("Cannot support {}".format(mask_nonlinear))
def __init__(self,
num_features,
bottleneck_channels=32,
hidden_channels=128,
chunk_size=100,
hop_size=None,
num_blocks=6,
num_heads=4,
causal=True,
norm=True,
mask_nonlinear='relu',
n_sources=2,
eps=EPS):
super().__init__()
if hop_size is None:
hop_size = chunk_size // 2
self.num_features, self.n_sources = num_features, n_sources
self.chunk_size, self.hop_size = chunk_size, hop_size
self.bottleneck_conv1d = nn.Conv1d(num_features,
bottleneck_channels,
kernel_size=1,
stride=1)
self.segment1d = Segment1d(chunk_size, hop_size)
self.norm2d = choose_layer_norm(bottleneck_channels,
causal=causal,
eps=eps)
self.dptransformer = DualPathTransformer(bottleneck_channels,
hidden_channels,
num_blocks=num_blocks,
num_heads=num_heads,
causal=causal,
norm=norm,
eps=eps)
self.overlap_add1d = OverlapAdd1d(chunk_size, hop_size)
self.gtu = GTU1d(bottleneck_channels, n_sources * num_features)
if mask_nonlinear == 'relu':
self.mask_nonlinear = nn.ReLU()
elif mask_nonlinear == 'sigmoid':
self.mask_nonlinear = nn.Sigmoid()
elif mask_nonlinear == 'softmax':
self.mask_nonlinear = nn.Softmax(dim=1)
else:
raise ValueError("Cannot support {}".format(mask_nonlinear))
def __init__(self,
in_channels,
out_channels=256,
skip_channels=256,
kernel_size=3,
stride=2,
dilation=1,
causal=True,
nonlinear=None,
norm=True,
dual_head=True,
eps=EPS):
super().__init__()
self.dual_head = dual_head
self.norm = norm
self.eps = eps
self.depthwise_conv1d = nn.Conv1d(in_channels,
in_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
groups=in_channels)
if nonlinear is not None:
if nonlinear == 'prelu':
self.nonlinear1d = nn.PReLU()
else:
raise ValueError("Not support {}".format(nonlinear))
self.nonlinear = True
else:
self.nonlinear = False
if norm:
self.norm1d = choose_layer_norm(in_channels, causal, eps=self.eps)
if dual_head:
self.output_pointwise_conv1d = nn.Conv1d(in_channels,
out_channels,
kernel_size=1,
stride=1)
self.skip_pointwise_conv1d = nn.Conv1d(in_channels,
skip_channels,
kernel_size=1,
stride=1)
def __init__(self, num_features, hidden_channels, causal, eps=EPS):
super().__init__()
self.num_features, self.hidden_channels = num_features, hidden_channels
if causal: # uni-direction
self.rnn = nn.LSTM(num_features,
hidden_channels,
batch_first=True,
bidirectional=False)
else: # bi-direction
self.rnn = nn.LSTM(num_features,
hidden_channels // 2,
batch_first=True,
bidirectional=True)
self.fc = nn.Linear(hidden_channels, num_features)
self.norm1d = choose_layer_norm(num_features, causal=False, eps=eps)