def __init__(self, channels, kernel_size, num_layers, dropout_p=0.0) -> torch.tensor: """Dilated Depth-wise Separable Convolution module. :: x |-> DDSConv(x) -> LayerNorm(x) -> GeLU(x) -> Conv1x1(x) -> LayerNorm(x) -> GeLU(x) -> + -> o |-------------------------------------------------------------------------------------^ Args: channels ([type]): [description] kernel_size ([type]): [description] num_layers ([type]): [description] dropout_p (float, optional): [description]. Defaults to 0.0. Returns: torch.tensor: Network output masked by the input sequence mask. """ super().__init__() self.num_layers = num_layers self.convs_sep = nn.ModuleList() self.convs_1x1 = nn.ModuleList() self.norms_1 = nn.ModuleList() self.norms_2 = nn.ModuleList() for i in range(num_layers): dilation = kernel_size ** i padding = (kernel_size * dilation - dilation) // 2 self.convs_sep.append( nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding) ) self.convs_1x1.append(nn.Conv1d(channels, channels, 1)) self.norms_1.append(LayerNorm2(channels)) self.norms_2.append(LayerNorm2(channels)) self.dropout = nn.Dropout(dropout_p)
def __init__( self, in_channels: int, out_channels: int, hidden_channels: int, hidden_channels_ffn: int, num_heads: int, num_layers: int, kernel_size=1, dropout_p=0.0, rel_attn_window_size: int = None, input_length: int = None, layer_norm_type: str = "1", ): super().__init__() self.hidden_channels = hidden_channels self.hidden_channels_ffn = hidden_channels_ffn self.num_heads = num_heads self.num_layers = num_layers self.kernel_size = kernel_size self.dropout_p = dropout_p self.rel_attn_window_size = rel_attn_window_size self.dropout = nn.Dropout(dropout_p) self.attn_layers = nn.ModuleList() self.norm_layers_1 = nn.ModuleList() self.ffn_layers = nn.ModuleList() self.norm_layers_2 = nn.ModuleList() for idx in range(self.num_layers): self.attn_layers.append( RelativePositionMultiHeadAttention( hidden_channels if idx != 0 else in_channels, hidden_channels, num_heads, rel_attn_window_size=rel_attn_window_size, dropout_p=dropout_p, input_length=input_length, )) if layer_norm_type == "1": self.norm_layers_1.append(LayerNorm(hidden_channels)) elif layer_norm_type == "2": self.norm_layers_1.append(LayerNorm2(hidden_channels)) else: raise ValueError(" [!] Unknown layer norm type") if hidden_channels != out_channels and (idx + 1) == self.num_layers: self.proj = nn.Conv1d(hidden_channels, out_channels, 1) self.ffn_layers.append( FeedForwardNetwork( hidden_channels, hidden_channels if (idx + 1) != self.num_layers else out_channels, hidden_channels_ffn, kernel_size, dropout_p=dropout_p, )) if layer_norm_type == "1": self.norm_layers_2.append( LayerNorm(hidden_channels if ( idx + 1) != self.num_layers else out_channels)) elif layer_norm_type == "2": self.norm_layers_2.append( LayerNorm2(hidden_channels if ( idx + 1) != self.num_layers else out_channels)) else: raise ValueError(" [!] Unknown layer norm type")