|
|
|
|
|
|
|
|
import mlx.core as mx |
|
|
from mlx.nn.layers.base import Module |
|
|
|
|
|
|
|
|
class Dropout(Module): |
|
|
r"""Randomly zero a portion of the elements during training. |
|
|
|
|
|
The remaining elements are multiplied with :math:`\frac{1}{1-p}` where |
|
|
:math:`p` is the probability of zeroing an element. This is done so the |
|
|
expected value of a given element will remain the same. |
|
|
|
|
|
Args: |
|
|
p (float): The probability to zero an element |
|
|
""" |
|
|
|
|
|
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 self._p_1 == 1 or not self.training: |
|
|
return x |
|
|
|
|
|
mask = mx.random.bernoulli(self._p_1, x.shape) |
|
|
|
|
|
return (1 / self._p_1) * mask * x |
|
|
|
|
|
|
|
|
class Dropout2d(Module): |
|
|
r"""Apply 2D 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 |
|
|
``NWHC`` or ``WHC`` where:``N`` is the batch dimension,``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 |
|
|
beneficial for early convolution layers where adjacent pixels are |
|
|
correlated. In such case, traditional dropout may not effectively |
|
|
regularize activations. For more details, see [1]. |
|
|
|
|
|
[1]: Thompson, J., Goroshin, R., Jain, A., LeCun, Y. and Bregler C., 2015. |
|
|
Efficient Object Localization Using Convolutional Networks. CVPR 2015. |
|
|
|
|
|
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 (3, 4): |
|
|
raise ValueError( |
|
|
f"Received input with {x.ndim} dimensions. Expected 3 or 4 dimensions." |
|
|
) |
|
|
|
|
|
if self._p_1 == 1 or not self.training: |
|
|
return x |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
mask_shape = x.shape |
|
|
mask_shape[-2] = mask_shape[-3] = 1 |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
|
