Explicit padding for ConformerConvolutionV1 (#16)

Author: JackTemaki

i6_models/parts/conformer/convolution.py

@@ -23,25 +23,35 @@ class ConformerConvolutionV1Config(ModelConfiguration):
     norm: Union[nn.Module, Callable[[torch.Tensor], torch.Tensor]]
     """normalization layer with input of shape [N,C,T]"""
 
+    def check_valid(self):
+        assert self.kernel_size % 2 == 1, "ConformerConvolutionV1 only supports odd kernel sizes"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.check_valid()
+
 
 class ConformerConvolutionV1(nn.Module):
     """
     Conformer convolution module.
     see also: https://github.com/espnet/espnet/blob/713e784c0815ebba2053131307db5f00af5159ea/espnet/nets/pytorch_backend/conformer/convolution.py#L13
+
+    Uses explicit padding for ONNX exportability, see:
+    https://github.com/pytorch/pytorch/issues/68880
     """
 
     def __init__(self, model_cfg: ConformerConvolutionV1Config):
         """
         :param model_cfg: model configuration for this module
         """
         super().__init__()
-
+        model_cfg.check_valid()
         self.pointwise_conv1 = nn.Linear(in_features=model_cfg.channels, out_features=2 * model_cfg.channels)
         self.depthwise_conv = nn.Conv1d(
             in_channels=model_cfg.channels,
             out_channels=model_cfg.channels,
             kernel_size=model_cfg.kernel_size,
-            padding="same",
+            padding=(model_cfg.kernel_size - 1) // 2,
             groups=model_cfg.channels,
         )
         self.pointwise_conv2 = nn.Linear(in_features=model_cfg.channels, out_features=model_cfg.channels)

tests/test_conformer.py

@@ -32,7 +32,6 @@ def get_output_shape(batch, time, features, norm=None, kernel_size=31, dropout=0
     assert get_output_shape(10, 1, 50) == (10, 1, 50)  # time dim 1
     assert get_output_shape(10, 10, 20, dropout=0.0) == (10, 10, 20)  # dropout 0
     assert get_output_shape(10, 10, 20, kernel_size=3) == (10, 10, 20)  # odd kernel size
-    assert get_output_shape(10, 10, 20, kernel_size=32) == (10, 10, 20)  # even kernel size
 
 
 def test_ConformerPositionwiseFeedForwardV1():