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

add dilation for conv 3d layers + test for 3d conv w/ dilation (#1802)

Browse Source
This commit is contained in:
Awni Hannun
2025-01-28 06:17:07 -08:00
committed by GitHub
parent ccb61d7aae
commit 1017ac4a9e
4 changed files with 22 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -179,6 +179,7 @@ class Conv3d(Module):
kernel_size (int or tuple): The size of the convolution filters.
stride (int or tuple, optional): The size of the stride when
applying the filter. Default: ``1``.
dilation (int or tuple, optional): The dilation of the convolution.
padding (int or tuple, optional): How many positions to 0-pad
the input with. Default: ``0``.
bias (bool, optional): If ``True`` add a learnable bias to the
@@ -192,6 +193,7 @@ class Conv3d(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__()
@@ -213,16 +215,18 @@ class Conv3d(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:3]}, 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.conv3d(x, self.weight, self.stride, self.padding)
y = mx.conv3d(x, self.weight, self.stride, self.padding, self.dilation)
if "bias" in self:
y = y + self.bias
return y

10
python/mlx/nn/layers/convolution_transpose.py
View File

@@ -159,6 +159,7 @@ class ConvTranspose3d(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``
"""
@@ -170,6 +171,7 @@ class ConvTranspose3d(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__()
@@ -191,16 +193,20 @@ class ConvTranspose3d(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:3]}, 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.conv_transpose3d(x, self.weight, self.stride, self.padding)
y = mx.conv_transpose3d(
x, self.weight, self.stride, self.padding, self.dilation
)
if "bias" in self:
y = y + self.bias
return y