[Ready]padding zeros

egs/aishell/ASR/conformer_ctc/conformer.py

@@ -248,7 +248,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -879,11 +881,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -897,6 +904,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         x = self.activation(self.norm(x))
 

egs/aishell/ASR/conformer_mmi/conformer.py

@@ -248,7 +248,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -879,11 +881,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -897,6 +904,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         x = self.activation(self.norm(x))
 

egs/aishell/ASR/transducer_stateless/conformer.py

@@ -246,7 +246,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -877,11 +879,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -895,6 +902,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         # x is (batch, channels, time)
         x = x.permute(0, 2, 1)

egs/gigaspeech/ASR/conformer_ctc/conformer.py

@@ -253,7 +253,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -890,11 +892,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -908,6 +915,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         if self.use_batchnorm:
             x = self.norm(x)

egs/librispeech/ASR/conformer_ctc/conformer.py

@@ -253,7 +253,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -890,11 +892,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -908,6 +915,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         if self.use_batchnorm:
             x = self.norm(x)

egs/librispeech/ASR/conformer_ctc2/conformer.py

@@ -268,7 +268,9 @@ def forward(
         src = src + self.dropout(src_att)
 
         # convolution module
-        src = src + self.dropout(self.conv_module(src))
+        src = src + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
 
         # feed forward module
         src = src + self.dropout(self.feed_forward(src))
@@ -921,11 +923,16 @@ def __init__(
             initial_scale=0.25,
         )
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -941,6 +948,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
 
         x = self.deriv_balancer2(x)

egs/librispeech/ASR/conformer_mmi/conformer.py

@@ -247,7 +247,9 @@ def forward(
         residual = src
         if self.normalize_before:
             src = self.norm_conv(src)
-        src = residual + self.dropout(self.conv_module(src))
+        src = residual + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
         if not self.normalize_before:
             src = self.norm_conv(src)
 
@@ -878,11 +880,16 @@ def __init__(
         )
         self.activation = Swish()
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -896,6 +903,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
         x = self.activation(self.norm(x))
 

egs/librispeech/ASR/pruned_transducer_stateless2/conformer.py

@@ -527,7 +527,9 @@ def forward(
         src = src + self.dropout(src_att)
 
         # convolution module
-        conv, _ = self.conv_module(src)
+        conv, _ = self.conv_module(
+            src, src_key_padding_mask=src_key_padding_mask
+        )
         src = src + self.dropout(conv)
 
         # feed forward module
@@ -1457,6 +1459,7 @@ def forward(
         x: Tensor,
         cache: Optional[Tensor] = None,
         right_context: int = 0,
+        src_key_padding_mask: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor]:
         """Compute convolution module.
 
@@ -1467,6 +1470,7 @@ def forward(
             right_context:
               How many future frames the attention can see in current chunk.
               Note: It's not that each individual frame has `right_context` frames
+            src_key_padding_mask: the mask for the src keys per batch (optional).
               of right context, some have more.
 
         Returns:
@@ -1486,6 +1490,8 @@ def forward(
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         if self.causal and self.lorder > 0:
             if cache is None:
                 # Make depthwise_conv causal by

egs/librispeech/ASR/pruned_transducer_stateless5/conformer.py

@@ -527,7 +527,9 @@ def forward(
         src = src + self.dropout(src_att)
 
         # convolution module
-        conv, _ = self.conv_module(src)
+        conv, _ = self.conv_module(
+            src, src_key_padding_mask=src_key_padding_mask
+        )
         src = src + self.dropout(conv)
 
         # feed forward module
@@ -1436,7 +1438,11 @@ def __init__(
         )
 
     def forward(
-        self, x: Tensor, cache: Optional[Tensor] = None, right_context: int = 0
+        self,
+        x: Tensor,
+        cache: Optional[Tensor] = None,
+        right_context: int = 0,
+        src_key_padding_mask: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor]:
         """Compute convolution module.
 
@@ -1448,6 +1454,7 @@ def forward(
               How many future frames the attention can see in current chunk.
               Note: It's not that each individual frame has `right_context` frames
               of right context, some have more.
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -1466,6 +1473,8 @@ def forward(
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         if self.causal and self.lorder > 0:
             if cache is None:
                 # Make depthwise_conv causal by

egs/librispeech/ASR/pruned_transducer_stateless6/conformer.py

@@ -264,7 +264,9 @@ def forward(
         src = src + self.dropout(src_att)
 
         # convolution module
-        src = src + self.dropout(self.conv_module(src))
+        src = src + self.dropout(
+            self.conv_module(src, src_key_padding_mask=src_key_padding_mask)
+        )
 
         # feed forward module
         src = src + self.dropout(self.feed_forward(src))
@@ -927,11 +929,16 @@ def __init__(
             initial_scale=0.25,
         )
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(
+        self,
+        x: Tensor,
+        src_key_padding_mask: Optional[Tensor] = None,
+    ) -> Tensor:
         """Compute convolution module.
 
         Args:
             x: Input tensor (#time, batch, channels).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
 
         Returns:
             Tensor: Output tensor (#time, batch, channels).
@@ -947,6 +954,8 @@ def forward(self, x: Tensor) -> Tensor:
         x = nn.functional.glu(x, dim=1)  # (batch, channels, time)
 
         # 1D Depthwise Conv
+        if src_key_padding_mask is not None:
+            x.masked_fill_(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
         x = self.depthwise_conv(x)
 
         x = self.deriv_balancer2(x)