zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 18:28:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Dilation for convolutional layers (#766)

Browse Source 
* add dilation parameter to Conv1d layer

* space here too

* add conv1d dilation test

* add dilation parameter for Conv2d layer

* conv2d dilation test
This commit is contained in:
Piotr Rybiec
2024-03-04 15:43:00 +01:00
committed by GitHub
parent bc06cb9ff6
commit 6a665ea6ed
2 changed files with 24 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/mlx/nn/layers/convolution.py
View File

@@ -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 @@ class Conv1d(Module):
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
bias: bool = True,
):
super().__init__()
@@ -48,17 +50,19 @@ class Conv1d(Module):
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
@@ -81,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``
"""
@@ -92,6 +97,7 @@ class Conv2d(Module):
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__()
@@ -111,16 +117,18 @@ class Conv2d(Module):
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