braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/modules/layers.py

@@ -0,0 +1,216 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+from __future__ import annotations
+
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.functional import drop_path
+
+
+class Ensure4d(nn.Module):
+    """Ensure the input tensor has 4 dimensions.
+
+    This is a small utility layer that repeatedly adds a singleton dimension at
+    the end until the input has shape ``(batch, channels, time, 1)``.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import Ensure4d
+    >>> module = Ensure4d()
+    >>> outputs = module(torch.randn(2, 3, 10))
+    >>> outputs.shape
+    torch.Size([2, 3, 10, 1])
+    """
+
+    def forward(self, x):
+        while len(x.shape) < 4:
+            x = x.unsqueeze(-1)
+        return x
+
+
+class Chomp1d(nn.Module):
+    """Remove samples from the end of a sequence.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import Chomp1d
+    >>> module = Chomp1d(chomp_size=5)
+    >>> inputs = torch.randn(2, 3, 20)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 3, 15])
+    """
+
+    def __init__(self, chomp_size):
+        super().__init__()
+        self.chomp_size = chomp_size
+
+    def extra_repr(self):
+        return "chomp_size={}".format(self.chomp_size)
+
+    def forward(self, x):
+        return x[:, :, : -self.chomp_size].contiguous()
+
+
+class TimeDistributed(nn.Module):
+    """Apply module on multiple windows.
+
+    Apply the provided module on a sequence of windows and return their
+    concatenation.
+    Useful with sequence-to-prediction models (e.g. sleep stager which must map
+    a sequence of consecutive windows to the label of the middle window in the
+    sequence).
+
+    Parameters
+    ----------
+    module : nn.Module
+        Module to be applied to the input windows. Must accept an input of
+        shape (batch_size, n_channels, n_times).
+
+    Examples
+    --------
+    >>> import torch
+    >>> from torch import nn
+    >>> from braindecode.modules import TimeDistributed
+    >>> module = TimeDistributed(nn.Conv1d(3, 4, kernel_size=3, padding=1))
+    >>> inputs = torch.randn(2, 5, 3, 20)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 5, 4])
+    """
+
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, x):
+        """
+        Parameters
+        ----------
+        x : torch.Tensor
+            Sequence of windows, of shape (batch_size, seq_len, n_channels,
+            n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Shape (batch_size, seq_len, output_size).
+        """
+        b, s, c, t = x.shape
+        out = self.module(x.view(b * s, c, t))
+        return out.view(b, s, -1)
+
+
+class DropPath(nn.Module):
+    """Drop paths, also known as Stochastic Depth, per sample.
+
+    When applied in main path of residual blocks.
+
+    Parameters
+    ----------
+    drop_prob: float (default=None)
+        Drop path probability (should be in range 0-1).
+
+    Notes
+    -----
+    Code copied and modified from VISSL facebookresearch:
+    https://github.com/facebookresearch/vissl/blob/0b5d6a94437bc00baed112ca90c9d78c6ccfbafb/vissl/models/model_helpers.py#L676
+
+    All rights reserved.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import DropPath
+    >>> module = DropPath(drop_prob=0.2)
+    >>> module.train()
+    >>> inputs = torch.randn(2, 3, 10)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 3, 10])
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+    # Utility function to print DropPath module
+    def extra_repr(self) -> str:
+        return f"p={self.drop_prob}"
+
+
+class SqueezeFinalOutput(nn.Module):
+    """
+
+    Removes empty dimension at end and potentially removes empty time
+    dimension. It does not just use squeeze as we never want to remove
+    first dimension.
+
+    Returns
+    -------
+    x: torch.Tensor
+        squeezed tensor
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self.squeeze = Rearrange("b c t 1 -> b c t")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 1) drop feature dim
+        x = self.squeeze(x)
+        # 2) drop time dim if singleton
+        if x.shape[-1] == 1:
+            x = x.squeeze(-1)
+        return x
+
+
+class SubjectLayers(nn.Module):
+    """Per-subject linear transformation layer.
+
+    Applies subject-specific linear transformations to the input. Each subject
+    owns an independent weight matrix, enabling personalized feature
+    processing.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        n_subjects: int,
+        init_id: bool = False,
+    ):
+        super().__init__()
+        self.weights = nn.Parameter(torch.randn(n_subjects, in_channels, out_channels))
+        if init_id:
+            if in_channels != out_channels:
+                raise AssertionError("init_id requires in_channels == out_channels")
+            self.weights.data[:] = torch.eye(in_channels)[None]
+        self.weights.data *= 1 / (in_channels**0.5)
+
+    def forward(self, x: torch.Tensor, subjects: torch.Tensor) -> torch.Tensor:
+        """Apply the subject-specific linear transforms."""
+        _, C, D = self.weights.shape
+        weights = self.weights.gather(0, subjects.view(-1, 1, 1).expand(-1, C, D))
+        return torch.einsum("bct,bcd->bdt", x, weights)
+
+    def __repr__(self) -> str:
+        S, C, D = self.weights.shape
+        return f"SubjectLayers({C}, {D}, {S})"

braindecode/modules/linear.py

@@ -0,0 +1,70 @@
+from torch import nn
+from torch.nn.utils.parametrize import register_parametrization
+
+from braindecode.modules.parametrization import MaxNorm, MaxNormParametrize
+
+
+class MaxNormLinear(nn.Linear):
+    """Linear layer with MaxNorm constraining on weights.
+
+    Equivalent of Keras tf.keras.Dense(..., kernel_constraint=max_norm())
+    [1]_ and [2]_. Implemented as advised in [3]_.
+
+    Parameters
+    ----------
+    in_features: int
+        Size of each input sample.
+    out_features: int
+        Size of each output sample.
+    bias: bool, optional
+        If set to ``False``, the layer will not learn an additive bias.
+        Default: ``True``.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import MaxNormLinear
+    >>> module = MaxNormLinear(10, 5, max_norm_val=2)
+    >>> inputs = torch.randn(2, 10)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 5])
+
+    References
+    ----------
+    .. [1] https://keras.io/api/layers/core_layers/dense/#dense-class
+    .. [2] https://www.tensorflow.org/api_docs/python/tf/keras/constraints/
+           MaxNorm
+    .. [3] https://discuss.pytorch.org/t/how-to-correctly-implement-in-place-
+           max-norm-constraint/96769
+    """
+
+    def __init__(
+        self, in_features, out_features, bias=True, max_norm_val=2, eps=1e-5, **kwargs
+    ):
+        super().__init__(
+            in_features=in_features, out_features=out_features, bias=bias, **kwargs
+        )
+        self._max_norm_val = max_norm_val
+        self._eps = eps
+        register_parametrization(self, "weight", MaxNorm(self._max_norm_val, self._eps))
+
+
+class LinearWithConstraint(nn.Linear):
+    """Linear layer with max-norm constraint on the weights.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import LinearWithConstraint
+    >>> module = LinearWithConstraint(10, 5, max_norm=1.0)
+    >>> inputs = torch.randn(2, 10)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 5])
+    """
+
+    def __init__(self, *args, max_norm=1.0, **kwargs):
+        super(LinearWithConstraint, self).__init__(*args, **kwargs)
+        self.max_norm = max_norm
+        register_parametrization(self, "weight", MaxNormParametrize(self.max_norm))

braindecode/modules/parametrization.py

@@ -0,0 +1,38 @@
+import torch
+from torch import nn
+
+
+class MaxNorm(nn.Module):
+    def __init__(self, max_norm_val=2.0, eps=1e-5):
+        super().__init__()
+        self.max_norm_val = max_norm_val
+        self.eps = eps
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        norm = X.norm(2, dim=0, keepdim=True)
+        denom = norm.clamp(min=self.max_norm_val / 2)
+        number = denom.clamp(max=self.max_norm_val)
+        return X * (number / (denom + self.eps))
+
+    def right_inverse(self, X: torch.Tensor) -> torch.Tensor:
+        # Assuming the forward scales X by a factor s,
+        # the right inverse would scale it back by 1/s.
+        norm = X.norm(2, dim=0, keepdim=True)
+        denom = norm.clamp(min=self.max_norm_val / 2)
+        number = denom.clamp(max=self.max_norm_val)
+        scale = number / (denom + self.eps)
+        return X / scale
+
+
+class MaxNormParametrize(nn.Module):
+    """
+    Enforce a max‑norm constraint on the rows of a weight tensor via parametrization.
+    """
+
+    def __init__(self, max_norm: float = 1.0):
+        super().__init__()
+        self.max_norm = max_norm
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        # Renormalize each "row" (dim=0 slice) to have at most self.max_norm L2-norm
+        return X.renorm(p=2, dim=0, maxnorm=self.max_norm)

braindecode/modules/stats.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Callable, Optional
+
+import torch
+from torch import nn
+
+
+class StatLayer(nn.Module):
+    """
+    Generic layer to compute a statistical function along a specified dimension.
+    Parameters
+    ----------
+    stat_fn : Callable
+        A function like torch.mean, torch.std, etc.
+    dim : int
+        Dimension along which to apply the function.
+    keepdim : bool, default=True
+        Whether to keep the reduced dimension.
+    clamp_range : tuple(float, float), optional
+        Used only for functions requiring clamping (e.g., log variance).
+    apply_log : bool, default=False
+        Whether to apply log after computation (used for LogVarLayer).
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import StatLayer
+    >>> module = StatLayer(stat_fn=torch.mean, dim=-1, keepdim=True)
+    >>> inputs = torch.randn(2, 3, 10)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 3, 1])
+    """
+
+    def __init__(
+        self,
+        stat_fn: Callable[..., torch.Tensor],
+        dim: int,
+        keepdim: bool = True,
+        clamp_range: Optional[tuple[float, float]] = None,
+        apply_log: bool = False,
+    ) -> None:
+        super().__init__()
+        self.stat_fn = stat_fn
+        self.dim = dim
+        self.keepdim = keepdim
+        self.clamp_range = clamp_range
+        self.apply_log = apply_log
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.stat_fn(x, dim=self.dim, keepdim=self.keepdim)
+        if self.clamp_range is not None:
+            out = torch.clamp(out, min=self.clamp_range[0], max=self.clamp_range[1])
+        if self.apply_log:
+            out = torch.log(out)
+        return out
+
+
+# make things more simple
+def _max_fn(x: torch.Tensor, dim: int, keepdim: bool) -> torch.Tensor:
+    return x.max(dim=dim, keepdim=keepdim)[0]
+
+
+def _power_fn(x: torch.Tensor, dim: int, keepdim: bool) -> torch.Tensor:
+    # compute mean of squared values along `dim`
+    return torch.mean(x**2, dim=dim, keepdim=keepdim)
+
+
+MeanLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.mean)
+MaxLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, _max_fn)
+VarLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.var)
+StdLayer: Callable[[int, bool], StatLayer] = partial(StatLayer, torch.std)
+LogVarLayer: Callable[[int, bool], StatLayer] = partial(
+    StatLayer,
+    torch.var,
+    clamp_range=(1e-6, 1e6),
+    apply_log=True,
+)
+
+LogPowerLayer: Callable[[int, bool], StatLayer] = partial(
+    StatLayer,
+    _power_fn,
+    clamp_range=(1e-4, 1e4),
+    apply_log=True,
+)

braindecode/modules/util.py

@@ -0,0 +1,85 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#          Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+import numpy as np
+from scipy.special import log_softmax
+
+
+def _pad_shift_array(x, stride=1):
+    """Zero-pad and shift rows of a 3D array.
+
+    E.g., used to align predictions of corresponding windows in
+    sequence-to-sequence models.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Array of shape (n_rows, n_classes, n_windows).
+    stride : int
+        Number of non-overlapping elements between two consecutive sequences.
+
+    Returns
+    -------
+    np.ndarray :
+        Array of shape (n_rows, n_classes, (n_rows - 1) * stride + n_windows)
+        where each row is obtained by zero-padding the corresponding row in
+        ``x`` before and after in the last dimension.
+    """
+    if x.ndim != 3:
+        raise NotImplementedError(
+            f"x must be of shape (n_rows, n_classes, n_windows), got {x.shape}"
+        )
+    x_padded = np.pad(x, ((0, 0), (0, 0), (0, (x.shape[0] - 1) * stride)))
+    orig_strides = x_padded.strides
+    new_strides = (
+        orig_strides[0] - stride * orig_strides[2],
+        orig_strides[1],
+        orig_strides[2],
+    )
+    return np.lib.stride_tricks.as_strided(x_padded, strides=new_strides)
+
+
+def aggregate_probas(logits, n_windows_stride=1):
+    """Aggregate predicted probabilities with self-ensembling.
+
+    Aggregate window-wise predicted probabilities obtained on overlapping
+    sequences of windows using multiplicative voting as described in
+    [Phan2018]_.
+
+    Parameters
+    ----------
+    logits : np.ndarray
+        Array of shape (n_sequences, n_classes, n_windows) containing the
+        logits (i.e. the raw unnormalized scores for each class) for each
+        window of each sequence.
+    n_windows_stride : int
+        Number of windows between two consecutive sequences. Default is 1
+        (maximally overlapping sequences).
+
+    Returns
+    -------
+    np.ndarray :
+        Array of shape ((n_rows - 1) * stride + n_windows, n_classes)
+        containing the aggregated predicted probabilities for each window
+        contained in the input sequences.
+
+    References
+    ----------
+    .. [Phan2018] Phan, H., Andreotti, F., Cooray, N., Chén, O. Y., &
+        De Vos, M. (2018). Joint classification and prediction CNN framework
+        for automatic sleep stage classification. IEEE Transactions on
+        Biomedical Engineering, 66(5), 1285-1296.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from braindecode.modules import aggregate_probas
+    >>> logits = np.random.randn(3, 4, 5)  # (n_sequences, n_classes, n_windows)
+    >>> probas = aggregate_probas(logits, n_windows_stride=1)
+    >>> probas.shape
+    (7, 4)
+    """
+    log_probas = log_softmax(logits, axis=1)
+    return _pad_shift_array(log_probas, stride=n_windows_stride).sum(axis=0).T

braindecode/modules/wrapper.py

@@ -0,0 +1,90 @@
+import torch
+from torch import nn
+
+
+class Expression(nn.Module):
+    """Compute given expression on forward pass.
+
+    Parameters
+    ----------
+    expression_fn : callable
+        Should accept variable number of objects of type
+        `torch.autograd.Variable` to compute its output.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from braindecode.modules import Expression
+    >>> module = Expression(lambda x: x**2)
+    >>> inputs = torch.randn(2, 3)
+    >>> outputs = module(inputs)
+    >>> outputs.shape
+    torch.Size([2, 3])
+    """
+
+    def __init__(self, expression_fn):
+        super().__init__()
+        self.expression_fn = expression_fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.expression_fn(x)
+
+    def __repr__(self):
+        if hasattr(self.expression_fn, "func") and hasattr(
+            self.expression_fn, "kwargs"
+        ):
+            expression_str = "{:s} {:s}".format(
+                self.expression_fn.func.__name__, str(self.expression_fn.kwargs)
+            )
+        elif hasattr(self.expression_fn, "__name__"):
+            expression_str = self.expression_fn.__name__
+        else:
+            expression_str = repr(self.expression_fn)
+        return self.__class__.__name__ + "(expression=%s) " % expression_str
+
+
+class IntermediateOutputWrapper(nn.Module):
+    """Wraps network model such that outputs of intermediate layers can be returned.
+    forward() returns list of intermediate activations in a network during forward pass.
+
+    Parameters
+    ----------
+    to_select : list
+        list of module names for which activation should be returned
+    model : model object
+        network model
+
+    Examples
+    --------
+    >>> model = Deep4Net()
+    >>> select_modules = ['conv_spat','conv_2','conv_3','conv_4'] # Specify intermediate outputs
+    >>> model_pert = IntermediateOutputWrapper(select_modules,model) # Wrap model
+
+    >>> import torch
+    >>> base = torch.nn.Sequential(torch.nn.Linear(10, 8), torch.nn.ReLU(), torch.nn.Linear(8, 2))
+    >>> wrapped = IntermediateOutputWrapper(to_select=["0", "2"], model=base)
+    >>> outputs = wrapped(torch.randn(4, 10))
+    >>> len(outputs)
+    2
+    """
+
+    def __init__(self, to_select, model):
+        if not len(list(model.children())) == len(list(model.named_children())):
+            raise Exception("All modules in model need to have names!")
+
+        super().__init__()
+
+        modules_list = model.named_children()
+        for key, module in modules_list:
+            self.add_module(key, module)
+            self._modules[key].load_state_dict(module.state_dict())
+        self._to_select = to_select
+
+    def forward(self, x):
+        # Call modules individually and append activation to output if module is in to_select
+        o = []
+        for name, module in self._modules.items():
+            x = module(x)
+            if name in self._to_select:
+                o.append(x)
+        return o