File size: 7,757 Bytes
283ba4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | # Copyright © 2023 Apple Inc.
import math
from typing import Union
import mlx.core as mx
from mlx.nn.layers.base import Module
class Conv1d(Module):
"""Applies a 1-dimensional convolution over the multi-channel input sequence.
The channels are expected to be last i.e. the input shape should be ``NLC`` where:
* ``N`` is the batch dimension
* ``L`` is the sequence length
* ``C`` is the number of input channels
Args:
in_channels (int): The number of input channels
out_channels (int): The number of output channels
kernel_size (int): The size of the convolution filters
stride (int, optional): The stride when applying the filter.
Default: ``1``.
padding (int, optional): How many positions to 0-pad the input with.
Default: ``0``.
dilation (int, optional): The dilation of the convolution.
groups (int, optional): The number of groups for the convolution.
Default: ``1``.
bias (bool, optional): If ``True`` add a learnable bias to the output.
Default: ``True``
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
):
super().__init__()
if in_channels % groups != 0:
raise ValueError(
f"The number of input channels ({in_channels}) must be "
f"divisible by the number of groups ({groups})"
)
scale = math.sqrt(1 / (in_channels * kernel_size))
self.weight = mx.random.uniform(
low=-scale,
high=scale,
shape=(out_channels, kernel_size, in_channels // groups),
)
if bias:
self.bias = mx.zeros((out_channels,))
self.padding = padding
self.dilation = dilation
self.stride = stride
self.groups = groups
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"groups={self.groups}, "
f"bias={'bias' in self}"
)
def __call__(self, x):
y = mx.conv1d(
x, self.weight, self.stride, self.padding, self.dilation, self.groups
)
if "bias" in self:
y = y + self.bias
return y
class Conv2d(Module):
"""Applies a 2-dimensional convolution over the multi-channel input image.
The channels are expected to be last i.e. the input shape should be ``NHWC`` where:
* ``N`` is the batch dimension
* ``H`` is the input image height
* ``W`` is the input image width
* ``C`` is the number of input channels
Args:
in_channels (int): The number of input channels.
out_channels (int): The number of output channels.
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``.
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.
groups (int, optional): The number of groups for the convolution.
Default: ``1``.
bias (bool, optional): If ``True`` add a learnable bias to the
output. Default: ``True``
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: Union[int, tuple],
stride: Union[int, tuple] = 1,
padding: Union[int, tuple] = 0,
dilation: Union[int, tuple] = 1,
groups: int = 1,
bias: bool = True,
):
super().__init__()
if in_channels % groups != 0:
raise ValueError(
f"The number of input channels ({in_channels}) must be "
f"divisible by the number of groups ({groups})"
)
kernel_size, stride, padding = map(
lambda x: (x, x) if isinstance(x, int) else x,
(kernel_size, stride, padding),
)
scale = math.sqrt(1 / (in_channels * kernel_size[0] * kernel_size[1]))
self.weight = mx.random.uniform(
low=-scale,
high=scale,
shape=(out_channels, *kernel_size, in_channels // groups),
)
if bias:
self.bias = mx.zeros((out_channels,))
self.padding = padding
self.stride = stride
self.dilation = dilation
self.groups = groups
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}, dilation={self.dilation}, "
f"groups={self.groups}, "
f"bias={'bias' in self}"
)
def __call__(self, x):
y = mx.conv2d(
x, self.weight, self.stride, self.padding, self.dilation, self.groups
)
if "bias" in self:
y = y + self.bias
return y
class Conv3d(Module):
"""Applies a 3-dimensional convolution over the multi-channel input image.
The channels are expected to be last i.e. the input shape should be ``NDHWC`` where:
* ``N`` is the batch dimension
* ``D`` is the input image depth
* ``H`` is the input image height
* ``W`` is the input image width
* ``C`` is the number of input channels
Args:
in_channels (int): The number of input channels.
out_channels (int): The number of output channels.
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
output. Default: ``True``
"""
def __init__(
self,
in_channels: int,
out_channels: int,
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__()
kernel_size, stride, padding = map(
lambda x: (x, x, x) if isinstance(x, int) else x,
(kernel_size, stride, padding),
)
scale = math.sqrt(
1 / (in_channels * kernel_size[0] * kernel_size[1] * kernel_size[2])
)
self.weight = mx.random.uniform(
low=-scale,
high=scale,
shape=(out_channels, *kernel_size, in_channels),
)
if bias:
self.bias = mx.zeros((out_channels,))
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:4]}, stride={self.stride}, "
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, self.dilation)
if "bias" in self:
y = y + self.bias
return y
|