braindecode 0.5.1__py3-none-any.whl

braindecode/models/sleep_stager_chambon_2018.py

@@ -0,0 +1,113 @@
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+
+import torch
+from torch import nn
+import numpy as np
+
+
+class SleepStagerChambon2018(nn.Module):
+    """
+    Sleep staging architecture from [1]_.
+
+    Convolutional neural network for sleep staging described in [1]_.
+
+    Parameters
+    ----------
+    n_channels : int
+        Number of EEG channels.
+    sfreq : float
+        EEG sampling frequency.
+    n_conv_chs : int
+        Number of convolutional channels. Set to 8 in [1]_.
+    time_conv_size_s : float
+        Size of filters in temporal convolution layers, in seconds. Set to 0.5
+        in [1]_ (64 samples at sfreq=128).
+    max_pool_size_s : float
+        Max pooling size, in seconds. Set to 0.125 in [1]_ (16 samples at
+        sfreq=128).
+    pad_size_s : float
+        Paddind size, in seconds. Set to 0.25 in [1]_ (half the temporal
+        convolution kernel size).
+    input_size_s : float
+        Size of the input, in seconds.
+    n_classes : int
+        Number of classes.
+    dropout : float
+        Dropout rate before the output dense layer.
+    apply_batch_norm : bool
+        If True, apply batch normalization after both temporal convolutional
+        layers.
+
+    References
+    ----------
+    .. [1] Chambon, S., Galtier, M. N., Arnal, P. J., Wainrib, G., &
+           Gramfort, A. (2018). A deep learning architecture for temporal sleep
+           stage classification using multivariate and multimodal time series.
+           IEEE Transactions on Neural Systems and Rehabilitation Engineering,
+           26(4), 758-769.
+    """
+    def __init__(self, n_channels, sfreq, n_conv_chs=8, time_conv_size_s=0.5,
+                 max_pool_size_s=0.125, pad_size_s=0.25, input_size_s=30,
+                 n_classes=5, dropout=0.25, apply_batch_norm=False):
+        super().__init__()
+
+        time_conv_size = np.ceil(time_conv_size_s * sfreq).astype(int)
+        max_pool_size = np.ceil(max_pool_size_s * sfreq).astype(int)
+        input_size = np.ceil(input_size_s * sfreq).astype(int)
+        pad_size = np.ceil(pad_size_s * sfreq).astype(int)
+
+        self.n_channels = n_channels
+
+        if n_channels > 1:
+            self.spatial_conv = nn.Conv2d(1, n_channels, (n_channels, 1))
+
+        batch_norm = nn.BatchNorm2d if apply_batch_norm else nn.Identity
+
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(
+                1, n_conv_chs, (1, time_conv_size), padding=(0, pad_size)),
+            batch_norm(n_conv_chs),
+            nn.ReLU(),
+            nn.MaxPool2d((1, max_pool_size)),
+            nn.Conv2d(
+                n_conv_chs, n_conv_chs, (1, time_conv_size),
+                padding=(0, pad_size)),
+            batch_norm(n_conv_chs),
+            nn.ReLU(),
+            nn.MaxPool2d((1, max_pool_size))
+        )
+        len_last_layer = self._len_last_layer(n_channels, input_size)
+        self.fc = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(len_last_layer, n_classes)
+        )
+
+    def _len_last_layer(self, n_channels, input_size):
+        self.feature_extractor.eval()
+        with torch.no_grad():
+            out = self.feature_extractor(
+                torch.Tensor(1, 1, n_channels, input_size))
+        self.feature_extractor.train()
+        return len(out.flatten())
+
+    def forward(self, x):
+        """
+        Forward pass.
+
+        Parameters
+        ---------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+        if x.ndim == 3:
+            x = x.unsqueeze(1)
+
+        if self.n_channels > 1:
+            x = self.spatial_conv(x)
+            x = x.transpose(1, 2)
+
+        x = self.feature_extractor(x)
+        return self.fc(x.flatten(start_dim=1))

braindecode/models/tcn.py

@@ -0,0 +1,143 @@
+# Authors: Patryk Chrabaszcz
+#          Lukas Gemein <l.gemein@gmail.com>
+#
+# License: BSD-3
+
+from torch import nn
+from torch.nn import init
+from torch.nn.utils import weight_norm
+
+from .modules import Ensure4d, Expression
+from .functions import squeeze_final_output
+
+
+class TCN(nn.Module):
+    """Temporal Convolutional Network (TCN) as described in [1]_.
+    Code adapted from https://github.com/locuslab/TCN/blob/master/TCN/tcn.py
+    Parameters
+    ----------
+    n_in_chans: int
+        number of input EEG channels
+    n_outputs: int
+        number of outputs of the decoding task (for example number of classes in
+        classification)
+    n_filters: int
+        number of output filters of each convolution
+    n_blocks: int
+        number of temporal blocks in the network
+    kernel_size: int
+        kernel size of the convolutions
+    drop_prob: float
+        dropout probability
+    add_log_softmax: bool
+        whether to add a log softmax layer
+    References
+    ----------
+    .. [1] Bai, S., Kolter, J. Z., & Koltun, V. (2018).
+       An empirical evaluation of generic convolutional and recurrent networks
+       for sequence modeling.
+       arXiv preprint arXiv:1803.01271.
+    """
+    def __init__(self, n_in_chans, n_outputs, n_blocks, n_filters, kernel_size,
+                 drop_prob, add_log_softmax):
+        super().__init__()
+        self.ensuredims = Ensure4d()
+        t_blocks = nn.Sequential()
+        for i in range(n_blocks):
+            n_inputs = n_in_chans if i == 0 else n_filters
+            dilation_size = 2 ** i
+            t_blocks.add_module("temporal_block_{:d}".format(i), TemporalBlock(
+                    n_inputs=n_inputs,
+                    n_outputs=n_filters,
+                    kernel_size=kernel_size,
+                    stride=1,
+                    dilation=dilation_size,
+                    padding=(kernel_size - 1) * dilation_size,
+                    drop_prob=drop_prob
+                ))
+        self.temporal_blocks = t_blocks
+        self.fc = nn.Linear(in_features=n_filters, out_features=n_outputs)
+        if add_log_softmax:
+            self.log_softmax = nn.LogSoftmax(dim=1)
+        self.squeeze = Expression(squeeze_final_output)
+
+        self.min_len = 1
+        for i in range(n_blocks):
+            dilation = 2 ** i
+            self.min_len += 2 * (kernel_size - 1) * dilation
+
+        # start in eval mode
+        self.eval()
+
+    def forward(self, x):
+        x = self.ensuredims(x)
+        # x is in format: B x C x T x 1
+        (batch_size, _, time_size, _) = x.size()
+        assert time_size >= self.min_len
+        # remove empty trailing dimension
+        x = x.squeeze(3)
+        x = self.temporal_blocks(x)
+        # Convert to: B x T x C
+        x = x.transpose(1, 2).contiguous()
+
+        fc_out = self.fc(x.view(batch_size * time_size, x.size(2)))
+        if hasattr(self, "log_softmax"):
+            fc_out = self.log_softmax(fc_out)
+        fc_out = fc_out.view(batch_size, time_size, fc_out.size(1))
+
+        out_size = 1 + max(0, time_size - self.min_len)
+        out = fc_out[:, -out_size:, :].transpose(1, 2)
+        # re-add 4th dimension for compatibility with braindecode
+        return self.squeeze(out[:, :, :, None])
+
+
+class TemporalBlock(nn.Module):
+    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation,
+                 padding, drop_prob):
+        super().__init__()
+        self.conv1 = weight_norm(nn.Conv1d(
+            n_inputs, n_outputs, kernel_size,
+            stride=stride, padding=padding, dilation=dilation))
+        self.chomp1 = Chomp1d(padding)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout2d(drop_prob)
+
+        self.conv2 = weight_norm(nn.Conv1d(
+            n_outputs, n_outputs, kernel_size,
+            stride=stride, padding=padding, dilation=dilation))
+        self.chomp2 = Chomp1d(padding)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout2d(drop_prob)
+
+        self.downsample = (nn.Conv1d(n_inputs, n_outputs, 1)
+                           if n_inputs != n_outputs else None)
+        self.relu = nn.ReLU()
+
+        init.normal_(self.conv1.weight, 0, 0.01)
+        init.normal_(self.conv2.weight, 0, 0.01)
+        if self.downsample is not None:
+            init.normal_(self.downsample.weight, 0, 0.01)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.chomp1(out)
+        out = self.relu1(out)
+        out = self.dropout1(out)
+        out = self.conv2(out)
+        out = self.chomp2(out)
+        out = self.relu2(out)
+        out = self.dropout2(out)
+        res = x if self.downsample is None else self.downsample(x)
+        return self.relu(out + res)
+
+
+class Chomp1d(nn.Module):
+    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()

braindecode/models/util.py

@@ -0,0 +1,66 @@
+import torch
+import numpy as np
+
+
+def to_dense_prediction_model(model, axis=(2, 3)):
+    """
+    Transform a sequential model with strides to a model that outputs
+    dense predictions by removing the strides and instead inserting dilations.
+    Modifies model in-place.
+
+    Parameters
+    ----------
+    model: torch.nn.Module
+        Model which modules will be modified
+    axis: int or (int,int)
+        Axis to transform (in terms of intermediate output axes)
+        can either be 2, 3, or (2,3).
+    
+    Notes
+    -----
+    Does not yet work correctly for average pooling.
+    Prior to version 0.1.7, there had been a bug that could move strides
+    backwards one layer.
+
+    """
+    if not hasattr(axis, "__len__"):
+        axis = [axis]
+    assert all([ax in [2, 3] for ax in axis]), "Only 2 and 3 allowed for axis"
+    axis = np.array(axis) - 2
+    stride_so_far = np.array([1, 1])
+    for module in model.modules():
+        if hasattr(module, "dilation"):
+            assert module.dilation == 1 or (module.dilation == (1, 1)), (
+                "Dilation should equal 1 before conversion, maybe the model is "
+                "already converted?"
+            )
+            new_dilation = [1, 1]
+            for ax in axis:
+                new_dilation[ax] = int(stride_so_far[ax])
+            module.dilation = tuple(new_dilation)
+        if hasattr(module, "stride"):
+            if not hasattr(module.stride, "__len__"):
+                module.stride = (module.stride, module.stride)
+            stride_so_far *= np.array(module.stride)
+            new_stride = list(module.stride)
+            for ax in axis:
+                new_stride[ax] = 1
+            module.stride = tuple(new_stride)
+
+
+def get_output_shape(model, in_chans, input_window_samples):
+    """Returns shape of neural network output for batch size equal 1.
+
+    Returns
+    -------
+    output_shape: tuple
+        shape of the network output for `batch_size==1` (1, ...)
+    """
+    with torch.no_grad():
+        dummy_input = torch.ones(
+            1, in_chans, input_window_samples,
+            dtype=next(model.parameters()).dtype,
+            device=next(model.parameters()).device,
+        )
+        output_shape = model(dummy_input).shape
+    return output_shape

braindecode/regressor.py

@@ -0,0 +1,240 @@
+import numpy as np
+from sklearn.metrics import get_scorer
+from skorch.callbacks import EpochTimer, BatchScoring, PrintLog, EpochScoring
+from skorch.classifier import NeuralNet
+from skorch.regressor import NeuralNetRegressor
+from skorch.utils import train_loss_score, valid_loss_score, noop, to_numpy
+
+from .training.scoring import PostEpochTrainScoring, CroppedTrialEpochScoring
+from .util import ThrowAwayIndexLoader, update_estimator_docstring
+
+
+class EEGRegressor(NeuralNetRegressor):
+    doc = """Regressor that calls loss function directly.
+
+    Parameters
+    ----------
+    cropped: bool (default=False)
+        Defines whether torch model passed to this class is cropped or not.
+        Currently used for callbacks definition.
+
+    callbacks: None or list of strings or list of Callback instances (default=None)
+        More callbacks, in addition to those returned by
+        ``get_default_callbacks``. Each callback should inherit from
+        :class:`skorch.callbacks.Callback`. If not ``None``, callbacks can be a
+        list of strings specifying `sklearn` scoring functions (for scoring
+        functions names see: https://scikit-learn.org/stable/modules/model_evaluation.html#scoring-parameter)
+        or a list of callbacks where the callback names are inferred from the
+        class name. Name conflicts are resolved by appending a count suffix
+        starting with 1, e.g. ``EpochScoring_1``. Alternatively,
+        a tuple ``(name, callback)`` can be passed, where ``name``
+        should be unique. Callbacks may or may not be instantiated.
+        The callback name can be used to set parameters on specific
+        callbacks (e.g., for the callback with name ``'print_log'``, use
+        ``net.set_params(callbacks__print_log__keys_ignored=['epoch',
+        'train_loss'])``).
+
+    iterator_train__shuffle: bool (default=True)
+        Defines whether train dataset will be shuffled. As skorch does not
+        shuffle the train dataset by default this one overwrites this option.
+
+    """
+    __doc__ = update_estimator_docstring(NeuralNetRegressor, doc)
+
+    def __init__(self, *args, cropped=False, callbacks=None,
+                 iterator_train__shuffle=True, **kwargs):
+        self.cropped = cropped
+        callbacks = self._parse_callbacks(callbacks)
+
+        super().__init__(*args,
+                         callbacks=callbacks,
+                         iterator_train__shuffle=iterator_train__shuffle,
+                         **kwargs)
+
+    def _parse_callbacks(self, callbacks):
+        callbacks_list = []
+        if callbacks is not None:
+            for callback in callbacks:
+                if isinstance(callback, tuple):
+                    callbacks_list.append(callback)
+                else:
+                    assert isinstance(callback, str)
+                    scoring = get_scorer(callback)
+                    scoring_name = scoring._score_func.__name__
+                    assert scoring_name.endswith(
+                        ('_score', '_error', '_deviance', '_loss'))
+                    if (scoring_name.endswith('_score') or
+                        callback.startswith('neg_')):
+                        lower_is_better = False
+                    else:
+                        lower_is_better = True
+                    train_name = f'train_{callback}'
+                    valid_name = f'valid_{callback}'
+                    if self.cropped:
+                        # In case of cropped decoding we are using braindecode
+                        # specific scoring created for cropped decoding
+                        train_scoring = CroppedTrialEpochScoring(
+                            callback, lower_is_better, on_train=True, name=train_name
+                        )
+                        valid_scoring = CroppedTrialEpochScoring(
+                            callback, lower_is_better, on_train=False, name=valid_name
+                        )
+                    else:
+                        train_scoring = PostEpochTrainScoring(
+                            callback, lower_is_better, name=train_name
+                        )
+                        valid_scoring = EpochScoring(
+                            callback, lower_is_better, on_train=False, name=valid_name
+                        )
+                    callbacks_list.extend([
+                        (train_name, train_scoring),
+                        (valid_name, valid_scoring)
+                    ])
+
+        return callbacks_list
+
+    # pylint: disable=arguments-differ
+    def get_loss(self, y_pred, y_true, *args, **kwargs):
+        """Return the loss for this batch by calling NeuralNet get_loss.
+        Parameters
+        ----------
+        y_pred : torch tensor
+          Predicted target values
+        y_true : torch tensor
+          True target values.
+        X : input data, compatible with skorch.dataset.Dataset
+          By default, you should be able to pass:
+            * numpy arrays
+            * torch tensors
+            * pandas DataFrame or Series
+            * scipy sparse CSR matrices
+            * a dictionary of the former three
+            * a list/tuple of the former three
+            * a Dataset
+          If this doesn't work with your data, you have to pass a
+          ``Dataset`` that can deal with the data.
+        training : bool (default=False)
+          Whether train mode should be used or not.
+
+        """
+        return NeuralNet.get_loss(self, y_pred, y_true, *args, **kwargs)
+
+    def get_iterator(self, dataset, training=False, drop_index=True):
+        iterator = super().get_iterator(dataset, training=training)
+        if drop_index:
+            return ThrowAwayIndexLoader(self, iterator, is_regression=True)
+        else:
+            return iterator
+
+    def on_batch_end(self, net, X, y, training=False, **kwargs):
+        # If training is false, assume that our loader has indices for this
+        # batch
+        if not training:
+            cbs = self._default_callbacks + self.callbacks
+            epoch_cbs = []
+            for name, cb in cbs:
+                if (cb.__class__.__name__ == 'CroppedTrialEpochScoring') and (
+                    hasattr(cb, 'window_inds_')) and (cb.on_train == False):
+                    epoch_cbs.append(cb)
+            # for trialwise decoding stuffs it might also be we don't have
+            # cropped loader, so no indices there
+            if len(epoch_cbs) > 0:
+                assert hasattr(self, '_last_window_inds')
+                for cb in epoch_cbs:
+                    cb.window_inds_.append(self._last_window_inds)
+                del self._last_window_inds
+
+    def predict_with_window_inds_and_ys(self, dataset):
+        preds = []
+        i_window_in_trials = []
+        i_window_stops = []
+        window_ys = []
+        for X, y, i in self.get_iterator(dataset, drop_index=False):
+            i_window_in_trials.append(i[0].cpu().numpy())
+            i_window_stops.append(i[2].cpu().numpy())
+            preds.append(to_numpy(self.forward(X)))
+            window_ys.append(y.cpu().numpy())
+        preds = np.concatenate(preds)
+        i_window_in_trials = np.concatenate(i_window_in_trials)
+        i_window_stops = np.concatenate(i_window_stops)
+        window_ys = np.concatenate(window_ys)
+        return dict(
+            preds=preds, i_window_in_trials=i_window_in_trials,
+            i_window_stops=i_window_stops, window_ys=window_ys)
+
+    # Removes default EpochScoring callback computing 'accuracy' to work properly
+    # with cropped decoding.
+    @property
+    def _default_callbacks(self):
+        return [
+            ("epoch_timer", EpochTimer()),
+            (
+                "train_loss",
+                BatchScoring(
+                    train_loss_score,
+                    name="train_loss",
+                    on_train=True,
+                    target_extractor=noop,
+                ),
+            ),
+            (
+                "valid_loss",
+                BatchScoring(
+                    valid_loss_score, name="valid_loss", target_extractor=noop,
+                ),
+            ),
+            ("print_log", PrintLog()),
+        ]
+
+    # Method added to fix `predict_proba` behavior to return proper values
+    # in cropped mode even if regressor does not return probabilities.
+    # We implement `predict_proba` because it exists in skorch.NeuralNetRegressor.
+    def predict_proba(self, X):
+        """Return the output of the module's forward method as a numpy
+        array. In case of cropped decoding returns averaged values for
+        each trial.
+
+        If the module's forward method returns multiple outputs as a
+        tuple, it is assumed that the first output contains the
+        relevant information and the other values are ignored.
+        If all values are relevant or module's output for each crop
+        is needed, consider using :func:`~skorch.NeuralNet.forward`
+        instead.
+
+        Parameters
+        ----------
+        X : input data, compatible with skorch.dataset.Dataset
+          By default, you should be able to pass:
+
+            * numpy arrays
+            * torch tensors
+            * pandas DataFrame or Series
+            * scipy sparse CSR matrices
+            * a dictionary of the former three
+            * a list/tuple of the former three
+            * a Dataset
+
+          If this doesn't work with your data, you have to pass a
+          ``Dataset`` that can deal with the data.
+
+        Warnings
+        --------
+        Regressors predict regression targets, so output of this method
+        can't be interpreted as probabilities. We advise you to use
+        `predict` method instead of `predict_proba`.
+
+        Returns
+        -------
+        y_proba : numpy ndarray
+
+        """
+        y_pred = super().predict_proba(X)
+        # Normally, we have to average the predictions across crops/timesteps
+        # to get one prediction per window/trial
+        # Predictions may be already averaged in CroppedTrialEpochScoring (y_pred.shape==2).
+        # However, when predictions are computed outside of CroppedTrialEpochScoring
+        # we have to average predictions, hence the check if len(y_pred.shape) == 3
+        if self.cropped and len(y_pred.shape) == 3:
+            return y_pred.mean(-1)
+        else:
+            return y_pred

braindecode/samplers/__init__.py

@@ -0,0 +1,5 @@
+"""Classes to sample examples.
+"""
+
+from .base import RecordingSampler
+from .ssl import RelativePositioningSampler

braindecode/samplers/base.py

@@ -0,0 +1,78 @@
+"""
+Sampler classes.
+"""
+
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+from torch.utils.data.sampler import Sampler
+from sklearn.utils import check_random_state
+
+
+class RecordingSampler(Sampler):
+    """Base sampler simplifying sampling from recordings.
+
+    Parameters
+    ----------
+    metadata : pd.DataFrame
+        DataFrame with at least one of {subject, session, run} columns for each
+        window in the BaseConcatDataset to sample examples from. Normally
+        obtained with `BaseConcatDataset.get_metadata()`. For instance,
+        `metadata.head()` might look like this:
+
+           i_window_in_trial  i_start_in_trial  i_stop_in_trial  target  subject    session    run
+        0                  0                 0              500      -1        4  session_T  run_0
+        1                  1               500             1000      -1        4  session_T  run_0
+        2                  2              1000             1500      -1        4  session_T  run_0
+        3                  3              1500             2000      -1        4  session_T  run_0
+        4                  4              2000             2500      -1        4  session_T  run_0
+
+    random_state : np.RandomState | int | None
+        Random state.
+
+    Attributes
+    ----------
+    info : pd.DataFrame
+        Series with MultiIndex index which contains the subject, session, run
+        and window indices information in an easily accessible structure for
+        quick sampling of windows.
+    """
+    def __init__(self, metadata, random_state=None):
+        self.metadata = metadata
+        self._init_info()
+        self.rng = check_random_state(random_state)
+
+    def _init_info(self):
+        keys = [k for k in ['subject', 'session', 'run']
+                if k in self.metadata.columns]
+        if not keys:
+            raise ValueError(
+                'metadata must contain at least one of the following columns: '
+                'subject, session or run.')
+
+        self.metadata = self.metadata.reset_index().rename(
+            columns={'index': 'window_index'})
+        self.info = self.metadata.reset_index().groupby(keys)[
+            ['index', 'i_start_in_trial']].agg(['unique'])
+        self.info.columns = self.info.columns.get_level_values(0)
+
+    def sample_recording(self):
+        """Return a random recording index.
+        """
+        return self.rng.choice(self.n_recordings)
+
+    def sample_window(self, rec_ind=None):
+        """Return a specific window.
+        """
+        if rec_ind is None:
+            rec_ind = self.sample_recording()
+        win_ind = self.rng.choice(self.info.iloc[rec_ind]['index'])
+        return win_ind, rec_ind
+
+    def __iter__(self):
+        raise NotImplementedError
+
+    @property
+    def n_recordings(self):
+        return self.info.shape[0]