| from torch import nn |
|
|
| from .normalization import LayerNorm |
|
|
|
|
| class GatedConvBlock(nn.Module): |
| """Gated convolutional block as in https://arxiv.org/pdf/1612.08083.pdf |
| Args: |
| in_out_channels (int): number of input/output channels. |
| kernel_size (int): convolution kernel size. |
| dropout_p (float): dropout rate. |
| """ |
|
|
| def __init__(self, in_out_channels, kernel_size, dropout_p, num_layers): |
| super().__init__() |
| |
| self.dropout_p = dropout_p |
| self.num_layers = num_layers |
| |
| self.conv_layers = nn.ModuleList() |
| self.norm_layers = nn.ModuleList() |
| self.layers = nn.ModuleList() |
| for _ in range(num_layers): |
| self.conv_layers += [nn.Conv1d(in_out_channels, 2 * in_out_channels, kernel_size, padding=kernel_size // 2)] |
| self.norm_layers += [LayerNorm(2 * in_out_channels)] |
|
|
| def forward(self, x, x_mask): |
| o = x |
| res = x |
| for idx in range(self.num_layers): |
| o = nn.functional.dropout(o, p=self.dropout_p, training=self.training) |
| o = self.conv_layers[idx](o * x_mask) |
| o = self.norm_layers[idx](o) |
| o = nn.functional.glu(o, dim=1) |
| o = res + o |
| res = o |
| return o |
|
|
