fix params naming

Author: mmz33

nn/conformer.py

@@ -14,10 +14,10 @@ class _PositionwiseFeedForward(nn.Module):
       FF -> Activation -> Dropout -> FF
   """
 
-  def __init__(self, d_model: int, d_ff: int, dropout: float, activation, l2: float = 0.0):
+  def __init__(self, dim_model: int, dim_ff: int, dropout: float, activation, l2: float = 0.0):
     """
-    :param d_model:
-    :param d_ff:
+    :param dim_model:
+    :param dim_ff:
     :param dropout:
     :param activation:
     :param l2:
@@ -27,8 +27,8 @@ def __init__(self, d_model: int, d_ff: int, dropout: float, activation, l2: floa
     self.dropout = dropout
     self.activation = activation
 
-    self.linear1 = nn.Linear(n_out=d_ff, l2=l2)
-    self.linear2 = nn.Linear(n_out=d_model, l2=l2)
+    self.linear1 = nn.Linear(n_out=dim_ff, l2=l2)
+    self.linear2 = nn.Linear(n_out=dim_model, l2=l2)
 
   def forward(self, inp: LayerRef) -> LayerRef:
     return self.linear2(nn.dropout(self.activation(self.linear1(inp)), dropout=self.dropout))
@@ -40,17 +40,17 @@ class _ConformerConvBlock(nn.Module):
       FF -> GLU -> depthwise conv -> BN -> Swish -> FF
   """
 
-  def __init__(self, d_model: int, kernel_size: Tuple[int], l2: float = 0.0):
+  def __init__(self, dim_model: int, kernel_size: Tuple[int], l2: float = 0.0):
     """
-    :param d_model:
+    :param dim_model:
     :param kernel_size:
     :param l2:
     """
     super().__init__()
 
-    self.positionwise_conv1 = nn.Linear(n_out=d_model * 2, l2=l2)
-    self.depthwise_conv = nn.Conv(n_out=d_model, filter_size=kernel_size, groups=d_model, l2=l2, padding='same')
-    self.positionwise_conv2 = nn.Linear(n_out=d_model, l2=l2)
+    self.positionwise_conv1 = nn.Linear(n_out=dim_model * 2, l2=l2)
+    self.depthwise_conv = nn.Conv(n_out=dim_model, filter_size=kernel_size, groups=dim_model, l2=l2, padding='same')
+    self.positionwise_conv2 = nn.Linear(n_out=dim_model, l2=l2)
 
   def forward(self, inp: LayerRef) -> LayerRef:
     x_conv1 = self.positionwise_conv1(inp)
@@ -68,15 +68,15 @@ class _ConformerConvSubsampleLayer(nn.Module):
   """
 
   def __init__(self, filter_sizes: List[Tuple[int, ...]], pool_sizes: Union[List[Tuple[int, ...]], None],
-      channel_sizes: List[int], l2: float = 0.0, dropout: float = 0.3, act: str = 'relu',
+      channel_sizes: List[int], l2: float = 0.0, dropout: float = 0.3, activation: str = 'relu',
       padding: str = 'same'):
     """
     :param filter_sizes:
     :param pool_sizes:
     :param channel_sizes:
     :param l2:
     :param dropout:
-    :param act:
+    :param activation:
     :param padding:
     """
     super().__init__()
@@ -87,7 +87,7 @@ def __init__(self, filter_sizes: List[Tuple[int, ...]], pool_sizes: Union[List[T
     self.conv_layers = nn.ModuleList()
     for filter_size, channel_size in zip(filter_sizes, channel_sizes):
       self.conv_layers.append(
-        nn.Conv(l2=l2, activation=act, filter_size=filter_size, n_out=channel_size, padding=padding))
+        nn.Conv(l2=l2, activation=activation, filter_size=filter_size, n_out=channel_size, padding=padding))
 
   def forward(self, inp: LayerRef) -> LayerRef:
     x = nn.split_dims(inp, axis='F', dims=(-1, 1))
@@ -106,31 +106,31 @@ class ConformerEncoderLayer(nn.Module):
   Represents a conformer block
   """
 
-  def __init__(self, conv_kernel_size: Tuple[int], ff_act, ff_dim: int, dropout: float, att_dropout: float,
-      enc_key_dim: int, att_n_heads: int, l2: float):
+  def __init__(self, conv_kernel_size: Tuple[int], activation_ff, dim_ff: int, dropout: float, att_dropout: float,
+      enc_key_dim: int, num_heads: int, l2: float):
     """
     :param conv_kernel_size:
-    :param ff_act:
+    :param activation_ff:
     :param ff_dim:
     :param dropout:
     :param att_dropout:
     :param enc_key_dim:
-    :param att_n_heads:
+    :param num_heads:
     :param l2:
     """
     super().__init__()
 
     self.dropout = dropout
 
     self.ffn1 = _PositionwiseFeedForward(
-      d_model=enc_key_dim, d_ff=ff_dim, dropout=dropout, activation=ff_act, l2=l2)
+      dim_model=enc_key_dim, dim_ff=dim_ff, dropout=dropout, activation=activation_ff, l2=l2)
 
     self.ffn2 = _PositionwiseFeedForward(
-      d_model=enc_key_dim, d_ff=ff_dim, dropout=dropout, activation=ff_act, l2=l2)
+      dim_model=enc_key_dim, dim_ff=dim_ff, dropout=dropout, activation=activation_ff, l2=l2)
 
-    self.conv_module = _ConformerConvBlock(d_model=enc_key_dim, kernel_size=conv_kernel_size)
+    self.conv_module = _ConformerConvBlock(dim_model=enc_key_dim, kernel_size=conv_kernel_size)
 
-    self.mhsa_module = MultiheadAttention(d_model, att_n_heads, dropout=att_dropout)  # TODO: to be implemented
+    self.mhsa_module = self.conv_module #MultiheadAttention(enc_key_dim, num_heads, dropout=att_dropout)  # TODO: to be implemented
 
   def forward(self, inp: LayerRef) -> LayerRef:
     # FFN
@@ -163,8 +163,8 @@ class ConformerEncoder(nn.Module):
   """
 
   def __init__(self, encoder_layer: nn.Module, num_blocks: int, conv_kernel_size: Tuple[int, ...] = (32,),
-      ff_act=nn.swish, ff_dim: int = 512, dropout: float = 0.1, att_dropout: float = 0.1, enc_key_dim: int = 256,
-      att_n_heads: int = 4, l2: float = 0.0):
+      activation_ff=nn.swish, dim_ff: int = 512, dropout: float = 0.1, att_dropout: float = 0.1, enc_key_dim: int = 256,
+      num_heads: int = 4, l2: float = 0.0):
     """
     :param encoder_layer:
     :param num_blocks:
@@ -189,8 +189,8 @@ def __init__(self, encoder_layer: nn.Module, num_blocks: int, conv_kernel_size:
 
     self.conformer_blocks = nn.Sequential([
       encoder_layer(
-        conv_kernel_size=conv_kernel_size, ff_act=ff_act, ff_dim=ff_dim, dropout=dropout,
-        att_dropout=att_dropout, enc_key_dim=enc_key_dim, att_n_heads=att_n_heads, l2=l2
+        conv_kernel_size=conv_kernel_size, activation_ff=activation_ff, dim_ff=dim_ff, dropout=dropout,
+        att_dropout=att_dropout, enc_key_dim=enc_key_dim, num_heads=num_heads, l2=l2
       )
       for _ in range(num_blocks)
     ])
@@ -200,4 +200,4 @@ def forward(self, inp: LayerRef) -> LayerRef:
     x_linear = self.linear(x_subsample)
     x = nn.dropout(x_linear, dropout=self.dropout)
     x = self.conformer_blocks(x)
-    return x
+    return x