From ad96f69c81a12cbe8c09de37994d51e4d4fa7e9f Mon Sep 17 00:00:00 2001 From: Piotrek Rybiec Date: Fri, 1 Mar 2024 13:10:54 +0100 Subject: [PATCH 1/5] add dilation parameter to Conv1d layer --- python/mlx/nn/layers/convolution.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/mlx/nn/layers/convolution.py b/python/mlx/nn/layers/convolution.py index c6928e188..d8e173782 100644 --- a/python/mlx/nn/layers/convolution.py +++ b/python/mlx/nn/layers/convolution.py @@ -23,6 +23,7 @@ class Conv1d(Module): Default: 1. padding (int, optional): How many positions to 0-pad the input with. Default: 0. + dilation (int, optional): The dilation of the convolution. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ @@ -34,6 +35,7 @@ def __init__( kernel_size: int, stride: int = 1, padding: int = 0, + dilation: int = 1, bias: bool = True, ): super().__init__() @@ -48,17 +50,19 @@ def __init__( self.bias = mx.zeros((out_channels,)) self.padding = padding + self.dilation = dilation self.stride = stride def _extra_repr(self): return ( f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " f"kernel_size={self.weight.shape[1]}, stride={self.stride}, " - f"padding={self.padding}, bias={'bias' in self}" + f"padding={self.padding}, dilation={self.dilation}," + f"bias={'bias' in self}" ) def __call__(self, x): - y = mx.conv1d(x, self.weight, self.stride, self.padding) + y = mx.conv1d(x, self.weight, self.stride, self.padding, self.dilation) if "bias" in self: y = y + self.bias return y From 7b12f2c7d2bbc5348a066f1d0dd125872a9eaf71 Mon Sep 17 00:00:00 2001 From: Piotrek Rybiec Date: Fri, 1 Mar 2024 14:51:21 +0100 Subject: [PATCH 2/5] space here too --- python/mlx/nn/layers/convolution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/mlx/nn/layers/convolution.py b/python/mlx/nn/layers/convolution.py index d8e173782..32a1fa1a8 100644 --- a/python/mlx/nn/layers/convolution.py +++ b/python/mlx/nn/layers/convolution.py @@ -57,7 +57,7 @@ def _extra_repr(self): return ( f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " f"kernel_size={self.weight.shape[1]}, stride={self.stride}, " - f"padding={self.padding}, dilation={self.dilation}," + f"padding={self.padding}, dilation={self.dilation}, " f"bias={'bias' in self}" ) From 87dc8cd8f5b60c70f732b6326a1f641589fd81b2 Mon Sep 17 00:00:00 2001 From: Piotrek Rybiec Date: Fri, 1 Mar 2024 15:37:56 +0100 Subject: [PATCH 3/5] add conv1d dilation test --- python/tests/test_nn.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 99154d3f6..bec6ff80f 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -586,6 +586,13 @@ def test_conv1d(self): self.assertEqual(y.shape, (N, (L - ks + 1) // 2, C_out)) self.assertTrue("bias" in c.parameters()) + dil = 2 + c = nn.Conv1d( + in_channels=C_in, out_channels=C_out, kernel_size=ks, dilation=dil + ) + y = c(x) + self.assertEqual(y.shape, (N, L - (ks - 1) * dil, C_out)) + c = nn.Conv1d(in_channels=C_in, out_channels=C_out, kernel_size=ks, bias=False) self.assertTrue("bias" not in c.parameters()) From 1262eea21305e8d10cb01c987875c364998ba343 Mon Sep 17 00:00:00 2001 From: Piotrek Rybiec Date: Fri, 1 Mar 2024 15:54:46 +0100 Subject: [PATCH 4/5] add dilation parameter for Conv2d layer --- python/mlx/nn/layers/convolution.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/mlx/nn/layers/convolution.py b/python/mlx/nn/layers/convolution.py index 32a1fa1a8..6e1c9780e 100644 --- a/python/mlx/nn/layers/convolution.py +++ b/python/mlx/nn/layers/convolution.py @@ -85,6 +85,7 @@ class Conv2d(Module): applying the filter. Default: 1. padding (int or tuple, optional): How many positions to 0-pad the input with. Default: 0. + dilation (int or tuple, optional): The dilation of the convolution. bias (bool, optional): If ``True`` add a learnable bias to the output. Default: ``True`` """ @@ -96,6 +97,7 @@ def __init__( kernel_size: Union[int, tuple], stride: Union[int, tuple] = 1, padding: Union[int, tuple] = 0, + dilation: Union[int, tuple] = 1, bias: bool = True, ): super().__init__() @@ -115,16 +117,18 @@ def __init__( self.padding = padding self.stride = stride + self.dilation = dilation def _extra_repr(self): return ( f"{self.weight.shape[-1]}, {self.weight.shape[0]}, " f"kernel_size={self.weight.shape[1:2]}, stride={self.stride}, " - f"padding={self.padding}, bias={'bias' in self}" + f"padding={self.padding}, dilation={self.dilation}, " + f"bias={'bias' in self}" ) def __call__(self, x): - y = mx.conv2d(x, self.weight, self.stride, self.padding) + y = mx.conv2d(x, self.weight, self.stride, self.padding, self.dilation) if "bias" in self: y = y + self.bias return y From 814ba9602a5e4dd047c4a49ac8fbe7f9dd2877f0 Mon Sep 17 00:00:00 2001 From: Piotrek Rybiec Date: Fri, 1 Mar 2024 17:13:19 +0100 Subject: [PATCH 5/5] conv2d dilation test --- python/tests/test_nn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index bec6ff80f..678acfd5b 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -639,6 +639,11 @@ def test_conv2d(self): self.assertEqual(y.shape, (4, 3, 3, 8)) self.assertLess(mx.abs(y - c.weight.sum((1, 2, 3))).max(), 1e-4) + c = nn.Conv2d(3, 8, 3, dilation=2) + y = c(x) + self.assertEqual(y.shape, (4, 4, 4, 8)) + self.assertLess(mx.abs(y - c.weight.sum((1, 2, 3))).max(), 1e-4) + def test_sequential(self): x = mx.ones((10, 2)) m = nn.Sequential(nn.Linear(2, 10), nn.ReLU(), nn.Linear(10, 1))