From 30df8c44b01195eb2984a4e666979c1ce98679d1 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Thu, 28 Dec 2023 23:06:59 +0530 Subject: [PATCH 1/6] feat: Add Dropout3d layer to nn.layers --- docs/src/python/nn/layers.rst | 1 + python/mlx/nn/layers/__init__.py | 2 +- python/mlx/nn/layers/dropout.py | 48 ++++++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/docs/src/python/nn/layers.rst b/docs/src/python/nn/layers.rst index aa59e0af24..7ead319fde 100644 --- a/docs/src/python/nn/layers.rst +++ b/docs/src/python/nn/layers.rst @@ -28,6 +28,7 @@ Layers GroupNorm Dropout Dropout2d + Dropout3d Transformer MultiHeadAttention ALiBi diff --git a/python/mlx/nn/layers/__init__.py b/python/mlx/nn/layers/__init__.py index 4dbe96eb6c..5d4c493953 100644 --- a/python/mlx/nn/layers/__init__.py +++ b/python/mlx/nn/layers/__init__.py @@ -35,7 +35,7 @@ from mlx.nn.layers.base import Module from mlx.nn.layers.containers import Sequential from mlx.nn.layers.convolution import Conv1d, Conv2d -from mlx.nn.layers.dropout import Dropout, Dropout2d +from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d from mlx.nn.layers.embedding import Embedding from mlx.nn.layers.linear import Linear from mlx.nn.layers.normalization import BatchNorm, GroupNorm, LayerNorm, RMSNorm diff --git a/python/mlx/nn/layers/dropout.py b/python/mlx/nn/layers/dropout.py index f656f8db36..09c85163b9 100644 --- a/python/mlx/nn/layers/dropout.py +++ b/python/mlx/nn/layers/dropout.py @@ -86,3 +86,51 @@ def __call__(self, x): mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape) return (1 / self._p_1) * mask * x + +class Dropout3d(Module): + r"""Apply 3D channel-wise dropout during training. + + Randomly zero out entire channels independently with probability :math:`p`. + This layer expects the channels to be last, i.e., the input shape should be + `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth, + `H` is the input image height, `W` is the input image width, and `C` is + the number of input channels. + + The remaining channels are scaled by :math:`\frac{1}{1-p}` to + maintain the expected value of each element. Unlike traditional dropout, + which zeros individual entries, this layer zeros entire channels. This is + often beneficial for convolutional layers processing 3D data, like in + medical imaging or video processing. + + Args: + p (float): Probability of zeroing a channel during training. + """ + + def __init__(self, p: float = 0.5): + super().__init__() + + if p < 0 or p >= 1: + raise ValueError(f"The dropout probability {p} is not in [0, 1)") + + self._p_1 = 1 - p + + def _extra_repr(self): + return f"p={1-self._p_1}" + + def __call__(self, x): + if x.ndim not in (4, 5): + raise ValueError( + f"Received input with {x.ndim} dimensions. Expected 4 or 5 dimensions." + ) + + if self._p_1 == 1 or not self.training: + return x + + # Dropout is applied on the whole channel + # 4D input: (1, 1, 1, C) + # 5D input: (B, 1, 1, 1, C) + mask_shape = list(x.shape) + mask_shape[-2] = mask_shape[-3] = mask_shape[-4] = 1 + + mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape) + return (1 / self._p_1) * mask * x From 8b47c9bd639c6c306b03cfeb378aa4e07efcd911 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Sat, 30 Dec 2023 01:33:11 +0530 Subject: [PATCH 2/6] acknowledgement --- ACKNOWLEDGMENTS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index ab3d5e1af1..acd0996761 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -7,6 +7,7 @@ with a short description of your contribution(s) below. For example: MLX was developed with contributions from the following individuals: +- Nripesh Niketan: Added `dropout3d` ops. - Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. - Juarez Bochi: Fixed bug in cross attention. - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example. From 20f508130e876a721ac1f10fc39966bfd69ae8c2 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Sat, 30 Dec 2023 01:55:14 +0530 Subject: [PATCH 3/6] Add dropout tests to test_nn.py --- python/tests/test_nn.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 204145c01c..a65893ee4a 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -791,7 +791,24 @@ def test_log_cosh_loss(self): targets = mx.zeros((2, 4)) loss = nn.losses.log_cosh_loss(inputs, targets, reduction="mean") self.assertAlmostEqual(loss.item(), 0.433781, places=6) - + + def test_dropout(self): + x = mx.ones((2, 4)) + y = nn.Dropout(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float32) + + def test_dropout2d(self): + x = mx.ones((2, 4, 4, 4)) + y = nn.Dropout2d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float32) + + def test_dropout3d(self): + x = mx.ones((2, 4, 4, 4, 4)) + y = nn.Dropout3d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float32) if __name__ == "__main__": unittest.main() From 9f1550fde9ffc0c894b5ecef135076c9139c83e7 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Sat, 30 Dec 2023 01:55:49 +0530 Subject: [PATCH 4/6] run pre-commit --- python/mlx/nn/layers/dropout.py | 5 +++-- python/tests/test_nn.py | 7 ++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/python/mlx/nn/layers/dropout.py b/python/mlx/nn/layers/dropout.py index 09c85163b9..18b9b03a6e 100644 --- a/python/mlx/nn/layers/dropout.py +++ b/python/mlx/nn/layers/dropout.py @@ -87,13 +87,14 @@ def __call__(self, x): mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape) return (1 / self._p_1) * mask * x + class Dropout3d(Module): r"""Apply 3D channel-wise dropout during training. Randomly zero out entire channels independently with probability :math:`p`. This layer expects the channels to be last, i.e., the input shape should be - `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth, - `H` is the input image height, `W` is the input image width, and `C` is + `NDHWC` or `DHWC` where: `N` is the batch dimension, `D` is the depth, + `H` is the input image height, `W` is the input image width, and `C` is the number of input channels. The remaining channels are scaled by :math:`\frac{1}{1-p}` to diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index a65893ee4a..34bd563d4c 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -791,24 +791,25 @@ def test_log_cosh_loss(self): targets = mx.zeros((2, 4)) loss = nn.losses.log_cosh_loss(inputs, targets, reduction="mean") self.assertAlmostEqual(loss.item(), 0.433781, places=6) - + def test_dropout(self): x = mx.ones((2, 4)) y = nn.Dropout(0.5)(x) self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) - + def test_dropout2d(self): x = mx.ones((2, 4, 4, 4)) y = nn.Dropout2d(0.5)(x) self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) - + def test_dropout3d(self): x = mx.ones((2, 4, 4, 4, 4)) y = nn.Dropout3d(0.5)(x) self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) + if __name__ == "__main__": unittest.main() From 9765b95a0d9799d068373cbc9ba4e5a9f9650b9c Mon Sep 17 00:00:00 2001 From: NripeshN Date: Mon, 1 Jan 2024 03:22:34 +0530 Subject: [PATCH 5/6] Add activation functions and dropout3d ops --- ACKNOWLEDGMENTS.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index acd0996761..398f136b27 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -7,8 +7,7 @@ with a short description of your contribution(s) below. For example: MLX was developed with contributions from the following individuals: -- Nripesh Niketan: Added `dropout3d` ops. -- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. +- Nripesh Niketan: Added `softsign`, `softmax`, `hardswish`, `logsoftmax` activation functions. Added `dropout3d` ops. - Juarez Bochi: Fixed bug in cross attention. - Justin Deschenaux: Sine, Cosine, arange, randint, truncated normal, bernoulli, lion optimizer, Dropout2d, linear and logistic regression python example. - Diogo Da Cruz: Added tri, tril, triu and safetensor support From c36fd76fbfd1305cef30dc8812c183e7b622b317 Mon Sep 17 00:00:00 2001 From: NripeshN Date: Mon, 1 Jan 2024 03:29:34 +0530 Subject: [PATCH 6/6] Add dropout tests for bfloat16 and float16 --- python/tests/test_nn.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/python/tests/test_nn.py b/python/tests/test_nn.py index 34bd563d4c..fa6b4bae38 100644 --- a/python/tests/test_nn.py +++ b/python/tests/test_nn.py @@ -798,18 +798,48 @@ def test_dropout(self): self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) + x = mx.ones((2, 4), dtype=mx.bfloat16) + y = nn.Dropout(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.bfloat16) + + x = mx.ones((2, 4), dtype=mx.float16) + y = nn.Dropout(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float16) + def test_dropout2d(self): x = mx.ones((2, 4, 4, 4)) y = nn.Dropout2d(0.5)(x) self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) + x = mx.ones((2, 4, 4, 4), dtype=mx.bfloat16) + y = nn.Dropout2d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.bfloat16) + + x = mx.ones((2, 4, 4, 4), dtype=mx.float16) + y = nn.Dropout2d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float16) + def test_dropout3d(self): x = mx.ones((2, 4, 4, 4, 4)) y = nn.Dropout3d(0.5)(x) self.assertTrue(y.shape, x.shape) self.assertTrue(y.dtype, mx.float32) + x = mx.ones((2, 4, 4, 4, 4), dtype=mx.bfloat16) + y = nn.Dropout3d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.bfloat16) + + x = mx.ones((2, 4, 4, 4, 4), dtype=mx.float16) + y = nn.Dropout3d(0.5)(x) + self.assertTrue(y.shape, x.shape) + self.assertTrue(y.dtype, mx.float16) + if __name__ == "__main__": unittest.main()