braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/samplers/ssl.py

@@ -0,0 +1,263 @@
+"""
+Self-supervised learning samplers.
+"""
+
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#          Young Truong <dt.young112@gmail.com>
+#
+# License: BSD (3-clause)
+
+import warnings
+
+import numpy as np
+import torch.distributed as dist
+
+from . import DistributedRecordingSampler, RecordingSampler
+
+
+class RelativePositioningSampler(RecordingSampler):
+    """Sample examples for the relative positioning task from [Banville2020]_.
+
+    Sample examples as tuples of two window indices, with a label indicating
+    whether the windows are close or far, as defined by tau_pos and tau_neg.
+
+    Parameters
+    ----------
+    metadata : pd.DataFrame
+        See RecordingSampler.
+    tau_pos : int
+        Size of the positive context, in samples. A positive pair contains two
+        windows x1 and x2 which are separated by at most `tau_pos` samples.
+    tau_neg : int
+        Size of the negative context, in samples. A negative pair contains two
+        windows x1 and x2 which are separated by at least `tau_neg` samples and
+        at most `tau_max` samples. Ignored if `same_rec_neg` is False.
+    n_examples : int
+        Number of pairs to extract.
+    tau_max : int | None
+        See `tau_neg`.
+    same_rec_neg : bool
+        If True, sample negative pairs from within the same recording. If
+        False, sample negative pairs from two different recordings.
+    random_state : None | np.RandomState | int
+        Random state.
+
+    References
+    ----------
+    .. [Banville2020] Banville, H., Chehab, O., Hyvärinen, A., Engemann, D. A.,
+           & Gramfort, A. (2020). Uncovering the structure of clinical EEG
+           signals with self-supervised learning.
+           arXiv preprint arXiv:2007.16104.
+    """
+
+    def __init__(
+        self,
+        metadata,
+        tau_pos,
+        tau_neg,
+        n_examples,
+        tau_max=None,
+        same_rec_neg=True,
+        random_state=None,
+    ):
+        super().__init__(metadata, random_state=random_state)
+
+        self.tau_pos = tau_pos
+        self.tau_neg = tau_neg
+        self.tau_max = np.inf if tau_max is None else tau_max
+        self.n_examples = n_examples
+        self.same_rec_neg = same_rec_neg
+
+        if not same_rec_neg and self.n_recordings < 2:
+            raise ValueError(
+                "More than one recording must be available when "
+                "using across-recording negative sampling."
+            )
+
+    def _sample_pair(self):
+        """Sample a pair of two windows."""
+        # Sample first window
+        win_ind1, rec_ind1 = self.sample_window()
+        ts1 = self.metadata.iloc[win_ind1]["i_start_in_trial"]
+        ts = self.info.iloc[rec_ind1]["i_start_in_trial"]
+
+        # Decide whether the pair will be positive or negative
+        pair_type = self.rng.binomial(1, 0.5)
+        win_ind2 = None
+        if pair_type == 0:  # Negative example
+            if self.same_rec_neg:
+                mask = ((ts <= ts1 - self.tau_neg) & (ts >= ts1 - self.tau_max)) | (
+                    (ts >= ts1 + self.tau_neg) & (ts <= ts1 + self.tau_max)
+                )
+            else:
+                rec_ind2 = rec_ind1
+                while rec_ind2 == rec_ind1:
+                    win_ind2, rec_ind2 = self.sample_window()
+        elif pair_type == 1:  # Positive example
+            mask = (ts >= ts1 - self.tau_pos) & (ts <= ts1 + self.tau_pos)
+
+        if win_ind2 is None:
+            mask[ts == ts1] = False  # same window cannot be sampled twice
+            if sum(mask) == 0:
+                raise NotImplementedError
+            win_ind2 = self.rng.choice(self.info.iloc[rec_ind1]["index"][mask])
+
+        return win_ind1, win_ind2, float(pair_type)
+
+    def presample(self):
+        """Presample examples.
+
+        Once presampled, the examples are the same from one epoch to another.
+        """
+        self.examples = [self._sample_pair() for _ in range(self.n_examples)]
+        return self
+
+    def __iter__(self):
+        """
+        Iterate over pairs.
+
+        Yields
+        ------
+        int
+            Position of the first window in the dataset.
+        int
+            Position of the second window in the dataset.
+        float
+            0 for a negative pair, 1 for a positive pair.
+        """
+        for i in range(self.n_examples):
+            if hasattr(self, "examples"):
+                yield self.examples[i]
+            else:
+                yield self._sample_pair()
+
+    def __len__(self):
+        return self.n_examples
+
+
+class DistributedRelativePositioningSampler(DistributedRecordingSampler):
+    """Sample examples for the relative positioning task from [Banville2020]_ in distributed mode.
+
+    Sample examples as tuples of two window indices, with a label indicating
+    whether the windows are close or far, as defined by tau_pos and tau_neg.
+
+    Parameters
+    ----------
+    metadata : pd.DataFrame
+        See RecordingSampler.
+    tau_pos : int
+        Size of the positive context, in samples. A positive pair contains two
+        windows x1 and x2 which are separated by at most `tau_pos` samples.
+    tau_neg : int
+        Size of the negative context, in samples. A negative pair contains two
+        windows x1 and x2 which are separated by at least `tau_neg` samples and
+        at most `tau_max` samples. Ignored if `same_rec_neg` is False.
+    n_examples : int
+        Number of pairs to extract.
+    tau_max : int | None
+        See `tau_neg`.
+    same_rec_neg : bool
+        If True, sample negative pairs from within the same recording. If
+        False, sample negative pairs from two different recordings.
+    random_state : None | np.RandomState | int
+        Random state.
+    kwargs: dict
+        Additional keyword arguments to pass to torch DistributedSampler.
+        See https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
+
+    References
+    ----------
+    .. [Banville2020] Banville, H., Chehab, O., Hyvärinen, A., Engemann, D. A.,
+        & Gramfort, A. (2020). Uncovering the structure of clinical EEG
+        signals with self-supervised learning.
+        arXiv preprint arXiv:2007.16104.
+    """
+
+    def __init__(
+        self,
+        metadata,
+        tau_pos,
+        tau_neg,
+        n_examples,
+        tau_max=None,
+        same_rec_neg=True,
+        random_state=None,
+        **kwargs,
+    ):
+        super().__init__(metadata, random_state=random_state, **kwargs)
+        self.tau_pos = tau_pos
+        self.tau_neg = tau_neg
+        self.tau_max = np.inf if tau_max is None else tau_max
+        self.same_rec_neg = same_rec_neg
+
+        self.n_examples = n_examples * self.n_recordings // self.info.shape[0]
+        warnings.warn(
+            f"Rank {dist.get_rank()} - Number of datasets: {self.n_recordings}"
+        )
+        warnings.warn(f"Rank {dist.get_rank()} - Number of samples: {self.n_examples}")
+
+        if not same_rec_neg and self.n_recordings < 2:
+            raise ValueError(
+                "More than one recording must be available when "
+                "using across-recording negative sampling."
+            )
+
+    def _sample_pair(self):
+        """Sample a pair of two windows."""
+        # Sample first window
+        win_ind1, rec_ind1 = self.sample_window()
+        ts1 = self.metadata.iloc[win_ind1]["i_start_in_trial"]
+        ts = self.info.iloc[rec_ind1]["i_start_in_trial"]
+
+        # Decide whether the pair will be positive or negative
+        pair_type = self.rng.binomial(1, 0.5)
+        win_ind2 = None
+        if pair_type == 0:  # Negative example
+            if self.same_rec_neg:
+                mask = ((ts <= ts1 - self.tau_neg) & (ts >= ts1 - self.tau_max)) | (
+                    (ts >= ts1 + self.tau_neg) & (ts <= ts1 + self.tau_max)
+                )
+            else:
+                rec_ind2 = rec_ind1
+                while rec_ind2 == rec_ind1:
+                    win_ind2, rec_ind2 = self.sample_window()
+        elif pair_type == 1:  # Positive example
+            mask = (ts >= ts1 - self.tau_pos) & (ts <= ts1 + self.tau_pos)
+
+        if win_ind2 is None:
+            mask[ts == ts1] = False  # same window cannot be sampled twice
+            if sum(mask) == 0:
+                raise NotImplementedError
+            win_ind2 = self.rng.choice(self.info.iloc[rec_ind1]["index"][mask])
+
+        return win_ind1, win_ind2, float(pair_type)
+
+    def presample(self):
+        """Presample examples.
+
+        Once presampled, the examples are the same from one epoch to another.
+        """
+        self.examples = [self._sample_pair() for _ in range(self.n_examples)]
+        return self
+
+    def __iter__(self):
+        """
+        Iterate over pairs.
+
+        Yields
+        ------
+        int
+            Position of the first window in the dataset.
+        int
+            Position of the second window in the dataset.
+        float
+            0 for a negative pair, 1 for a positive pair.
+        """
+        for i in range(self.n_examples):
+            if hasattr(self, "examples"):
+                yield self.examples[i]
+            else:
+                yield self._sample_pair()
+
+    def __len__(self):
+        return self.n_examples

braindecode/training/__init__.py

@@ -0,0 +1,23 @@
+"""
+Functionality for skorch-based training.
+"""
+
+from .losses import CroppedLoss, TimeSeriesLoss, mixup_criterion
+from .scoring import (
+    CroppedTimeSeriesEpochScoring,
+    CroppedTrialEpochScoring,
+    PostEpochTrainScoring,
+    predict_trials,
+    trial_preds_from_window_preds,
+)
+
+__all__ = [
+    "CroppedLoss",
+    "mixup_criterion",
+    "TimeSeriesLoss",
+    "CroppedTrialEpochScoring",
+    "PostEpochTrainScoring",
+    "CroppedTimeSeriesEpochScoring",
+    "trial_preds_from_window_preds",
+    "predict_trials",
+]

braindecode/training/callbacks.py

@@ -0,0 +1,23 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+from skorch.callbacks import Callback
+
+
+class MaxNormConstraintCallback(Callback):
+    def on_batch_end(self, net, training, *args, **kwargs):
+        if training:
+            model = net.module_
+            last_weight = None
+            for name, module in list(model.named_children()):
+                if hasattr(module, "weight") and (
+                    not module.__class__.__name__.startswith("BatchNorm")
+                ):
+                    module.weight.data = torch.renorm(
+                        module.weight.data, 2, 0, maxnorm=2
+                    )
+                    last_weight = module.weight
+            if last_weight is not None:
+                last_weight.data = torch.renorm(last_weight.data, 2, 0, maxnorm=0.5)

braindecode/training/losses.py

@@ -0,0 +1,105 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#          Maciej Sliwowski <maciek.sliwowski@gmail.com>
+#          Mohammed Fattouh <mo.fattouh@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+from torch import nn
+
+
+class CroppedLoss(nn.Module):
+    """Compute Loss after averaging predictions across time.
+    Assumes predictions are in shape:
+    n_batch size x n_classes x n_predictions (in time)"""
+
+    def __init__(self, loss_function):
+        super().__init__()
+        self.loss_function = loss_function
+
+    def forward(self, preds, targets):
+        """Forward pass.
+
+        Parameters
+        ----------
+        preds: torch.Tensor
+            Model's prediction with shape (batch_size, n_classes, n_times).
+        targets: torch.Tensor
+            Target labels with shape (batch_size, n_classes, n_times).
+        """
+        avg_preds = torch.mean(preds, dim=2)
+        avg_preds = avg_preds.squeeze(dim=1)
+        return self.loss_function(avg_preds, targets)
+
+
+class TimeSeriesLoss(nn.Module):
+    """Compute Loss between timeseries targets and predictions.
+    Assumes predictions are in shape:
+    n_batch size x n_classes x n_predictions (in time)
+    Assumes targets are in shape:
+    n_batch size x n_classes x window_len (in time)
+    If the targets contain NaNs, the NaNs will be masked out and the loss will be only computed for
+    predictions valid corresponding to valid target values."""
+
+    def __init__(self, loss_function):
+        super().__init__()
+        self.loss_function = loss_function
+
+    def forward(self, preds, targets):
+        """Forward pass.
+
+        Parameters
+        ----------
+        preds: torch.Tensor
+            Model's prediction with shape (batch_size, n_classes, n_times).
+        targets: torch.Tensor
+            Target labels with shape (batch_size, n_classes, n_times).
+        """
+        n_preds = preds.shape[-1]
+        # slice the targets to fit preds shape
+        targets = targets[:, :, -n_preds:]
+        # create valid targets mask
+        mask = ~torch.isnan(targets)
+        # select valid targets that have a matching predictions
+        masked_targets = targets[mask]
+        masked_preds = preds[mask]
+        return self.loss_function(masked_preds, masked_targets)
+
+
+def mixup_criterion(preds, target):
+    """Implements loss for Mixup for EEG data. See [1]_.
+
+    Implementation based on [2]_.
+
+    Parameters
+    ----------
+    preds : torch.Tensor
+        Predictions from the model.
+    target : torch.Tensor | list of torch.Tensor
+        For predictions without mixup, the targets as a tensor. If mixup has
+        been applied, a list containing the targets of the two mixed
+        samples and the mixing coefficients as tensors.
+
+    Returns
+    -------
+    loss : float
+        The loss value.
+
+    References
+    ----------
+    .. [1] Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, David Lopez-Paz
+       mixup: Beyond Empirical Risk Minimization
+       Online: https://arxiv.org/abs/1710.09412
+    .. [2] https://github.com/facebookresearch/mixup-cifar10/blob/master/train.py
+    """
+    if len(target) == 3:
+        # unpack target
+        y_a, y_b, lam = target
+        # compute loss per sample
+        loss_a = torch.nn.functional.nll_loss(preds, y_a, reduction="none")
+        loss_b = torch.nn.functional.nll_loss(preds, y_b, reduction="none")
+        # compute weighted mean
+        ret = torch.mul(lam, loss_a) + torch.mul(1 - lam, loss_b)
+        return ret.mean()
+    else:
+        return torch.nn.functional.nll_loss(preds, target)