|
|
import torch.nn as nn |
|
|
import torch |
|
|
|
|
|
|
|
|
class GLU(nn.Module): |
|
|
def __init__(self, input_num): |
|
|
super(GLU, self).__init__() |
|
|
self.sigmoid = nn.Sigmoid() |
|
|
self.linear = nn.Linear(input_num, input_num) |
|
|
|
|
|
def forward(self, x): |
|
|
lin = self.linear(x.permute(0, 2, 3, 1)) |
|
|
lin = lin.permute(0, 3, 1, 2) |
|
|
sig = self.sigmoid(x) |
|
|
res = lin * sig |
|
|
return res |
|
|
|
|
|
|
|
|
class ContextGating(nn.Module): |
|
|
def __init__(self, input_num): |
|
|
super(ContextGating, self).__init__() |
|
|
self.sigmoid = nn.Sigmoid() |
|
|
self.linear = nn.Linear(input_num, input_num) |
|
|
|
|
|
def forward(self, x): |
|
|
lin = self.linear(x.permute(0, 2, 3, 1)) |
|
|
lin = lin.permute(0, 3, 1, 2) |
|
|
sig = self.sigmoid(lin) |
|
|
res = x * sig |
|
|
return res |
|
|
|
|
|
|
|
|
class CNN(nn.Module): |
|
|
def __init__( |
|
|
self, |
|
|
n_in_channel, |
|
|
activation="Relu", |
|
|
conv_dropout=0, |
|
|
kernel_size=[3, 3, 3], |
|
|
padding=[1, 1, 1], |
|
|
stride=[1, 1, 1], |
|
|
nb_filters=[64, 64, 64], |
|
|
pooling=[(1, 4), (1, 4), (1, 4)], |
|
|
normalization="batch", |
|
|
**transformer_kwargs |
|
|
): |
|
|
""" |
|
|
Initialization of CNN network s |
|
|
|
|
|
Args: |
|
|
n_in_channel: int, number of input channel |
|
|
activation: str, activation function |
|
|
conv_dropout: float, dropout |
|
|
kernel_size: kernel size |
|
|
padding: padding |
|
|
stride: list, stride |
|
|
nb_filters: number of filters |
|
|
pooling: list of tuples, time and frequency pooling |
|
|
normalization: choose between "batch" for BatchNormalization and "layer" for LayerNormalization. |
|
|
""" |
|
|
super(CNN, self).__init__() |
|
|
|
|
|
self.nb_filters = nb_filters |
|
|
cnn = nn.Sequential() |
|
|
|
|
|
def conv(i, normalization="batch", dropout=None, activ="relu"): |
|
|
nIn = n_in_channel if i == 0 else nb_filters[i - 1] |
|
|
nOut = nb_filters[i] |
|
|
cnn.add_module( |
|
|
"conv{0}".format(i), |
|
|
nn.Conv2d(nIn, nOut, kernel_size[i], stride[i], padding[i]), |
|
|
) |
|
|
if normalization == "batch": |
|
|
cnn.add_module( |
|
|
"batchnorm{0}".format(i), |
|
|
nn.BatchNorm2d(nOut, eps=0.001, momentum=0.99), |
|
|
) |
|
|
elif normalization == "layer": |
|
|
cnn.add_module("layernorm{0}".format(i), nn.GroupNorm(1, nOut)) |
|
|
|
|
|
if activ.lower() == "leakyrelu": |
|
|
cnn.add_module("relu{0}".format(i), nn.LeakyReLU(0.2)) |
|
|
elif activ.lower() == "relu": |
|
|
cnn.add_module("relu{0}".format(i), nn.ReLU()) |
|
|
elif activ.lower() == "glu": |
|
|
cnn.add_module("glu{0}".format(i), GLU(nOut)) |
|
|
elif activ.lower() == "cg": |
|
|
cnn.add_module("cg{0}".format(i), ContextGating(nOut)) |
|
|
|
|
|
if dropout is not None: |
|
|
cnn.add_module("dropout{0}".format(i), nn.Dropout(dropout)) |
|
|
|
|
|
|
|
|
for i in range(len(nb_filters)): |
|
|
conv(i, normalization=normalization, dropout=conv_dropout, activ=activation) |
|
|
cnn.add_module( |
|
|
"pooling{0}".format(i), nn.AvgPool2d(pooling[i]) |
|
|
) |
|
|
|
|
|
self.cnn = cnn |
|
|
|
|
|
def forward(self, x): |
|
|
""" |
|
|
Forward step of the CNN module |
|
|
|
|
|
Args: |
|
|
x (Tensor): input batch of size (batch_size, n_channels, n_frames, n_freq) |
|
|
|
|
|
Returns: |
|
|
Tensor: batch embedded |
|
|
""" |
|
|
|
|
|
x = self.cnn(x) |
|
|
return x |
|
|
|
