braindecode 0.8__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/models/deep4.py

@@ -0,0 +1,322 @@
+# Authors: Robin Schirrmeister <robintibor@gmail.com>
+#
+# License: BSD (3-clause)
+
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.nn import init
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import (
+    AvgPool2dWithConv,
+    CombinedConv,
+    Ensure4d,
+    SqueezeFinalOutput,
+)
+
+
+class Deep4Net(EEGModuleMixin, nn.Sequential):
+    """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
+
+     .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
+        :align: center
+        :alt: CTNet Architecture
+
+    Model described in [Schirrmeister2017]_.
+
+    Parameters
+    ----------
+    final_conv_length: int | str
+        Length of the final convolution layer.
+        If set to "auto", n_times must not be None. Default: "auto".
+    n_filters_time: int
+        Number of temporal filters.
+    n_filters_spat: int
+        Number of spatial filters.
+    filter_time_length: int
+        Length of the temporal filter in layer 1.
+    pool_time_length: int
+        Length of temporal pooling filter.
+    pool_time_stride: int
+        Length of stride between temporal pooling filters.
+    n_filters_2: int
+        Number of temporal filters in layer 2.
+    filter_length_2: int
+        Length of the temporal filter in layer 2.
+    n_filters_3: int
+        Number of temporal filters in layer 3.
+    filter_length_3: int
+        Length of the temporal filter in layer 3.
+    n_filters_4: int
+        Number of temporal filters in layer 4.
+    filter_length_4: int
+        Length of the temporal filter in layer 4.
+    activation_first_conv_nonlin: nn.Module, default is nn.ELU
+        Non-linear activation function to be used after convolution in layer 1.
+    first_pool_mode: str
+        Pooling mode in layer 1. "max" or "mean".
+    first_pool_nonlin: callable
+        Non-linear activation function to be used after pooling in layer 1.
+    activation_later_conv_nonlin: nn.Module, default is nn.ELU
+        Non-linear activation function to be used after convolution in later layers.
+    later_pool_mode: str
+        Pooling mode in later layers. "max" or "mean".
+    later_pool_nonlin: callable
+        Non-linear activation function to be used after pooling in later layers.
+    drop_prob: float
+        Dropout probability.
+    split_first_layer: bool
+        Split first layer into temporal and spatial layers (True) or just use temporal (False).
+        There would be no non-linearity between the split layers.
+    batch_norm: bool
+        Whether to use batch normalisation.
+    batch_norm_alpha: float
+        Momentum for BatchNorm2d.
+    stride_before_pool: bool
+        Stride before pooling.
+
+
+    References
+    ----------
+    .. [Schirrmeister2017] Schirrmeister, R. T., Springenberg, J. T., Fiederer,
+       L. D. J., Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F.
+       & Ball, T. (2017).
+       Deep learning with convolutional neural networks for EEG decoding and
+       visualization.
+       Human Brain Mapping , Aug. 2017.
+       Online: http://dx.doi.org/10.1002/hbm.23730
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        final_conv_length="auto",
+        n_filters_time=25,
+        n_filters_spat=25,
+        filter_time_length=10,
+        pool_time_length=3,
+        pool_time_stride=3,
+        n_filters_2=50,
+        filter_length_2=10,
+        n_filters_3=100,
+        filter_length_3=10,
+        n_filters_4=200,
+        filter_length_4=10,
+        activation_first_conv_nonlin: nn.Module = nn.ELU,
+        first_pool_mode="max",
+        first_pool_nonlin: nn.Module = nn.Identity,
+        activation_later_conv_nonlin: nn.Module = nn.ELU,
+        later_pool_mode="max",
+        later_pool_nonlin: nn.Module = nn.Identity,
+        drop_prob=0.5,
+        split_first_layer=True,
+        batch_norm=True,
+        batch_norm_alpha=0.1,
+        stride_before_pool=False,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+
+        if final_conv_length == "auto":
+            assert self.n_times is not None
+        self.final_conv_length = final_conv_length
+        self.n_filters_time = n_filters_time
+        self.n_filters_spat = n_filters_spat
+        self.filter_time_length = filter_time_length
+        self.pool_time_length = pool_time_length
+        self.pool_time_stride = pool_time_stride
+        self.n_filters_2 = n_filters_2
+        self.filter_length_2 = filter_length_2
+        self.n_filters_3 = n_filters_3
+        self.filter_length_3 = filter_length_3
+        self.n_filters_4 = n_filters_4
+        self.filter_length_4 = filter_length_4
+        self.first_nonlin = activation_first_conv_nonlin
+        self.first_pool_mode = first_pool_mode
+        self.first_pool_nonlin = first_pool_nonlin
+        self.later_conv_nonlin = activation_later_conv_nonlin
+        self.later_pool_mode = later_pool_mode
+        self.later_pool_nonlin = later_pool_nonlin
+        self.drop_prob = drop_prob
+        self.split_first_layer = split_first_layer
+        self.batch_norm = batch_norm
+        self.batch_norm_alpha = batch_norm_alpha
+        self.stride_before_pool = stride_before_pool
+
+        # For the load_state_dict
+        # When padronize all layers,
+        # add the old's parameters here
+        self.mapping = {
+            "conv_time.weight": "conv_time_spat.conv_time.weight",
+            "conv_spat.weight": "conv_time_spat.conv_spat.weight",
+            "conv_time.bias": "conv_time_spat.conv_time.bias",
+            "conv_spat.bias": "conv_time_spat.conv_spat.bias",
+            "conv_classifier.weight": "final_layer.conv_classifier.weight",
+            "conv_classifier.bias": "final_layer.conv_classifier.bias",
+        }
+
+        if self.stride_before_pool:
+            conv_stride = self.pool_time_stride
+            pool_stride = 1
+        else:
+            conv_stride = 1
+            pool_stride = self.pool_time_stride
+        self.add_module("ensuredims", Ensure4d())
+        pool_class_dict = dict(max=nn.MaxPool2d, mean=AvgPool2dWithConv)
+        first_pool_class = pool_class_dict[self.first_pool_mode]
+        later_pool_class = pool_class_dict[self.later_pool_mode]
+        if self.split_first_layer:
+            self.add_module("dimshuffle", Rearrange("batch C T 1 -> batch 1 T C"))
+            self.add_module(
+                "conv_time_spat",
+                CombinedConv(
+                    in_chans=self.n_chans,
+                    n_filters_time=self.n_filters_time,
+                    n_filters_spat=self.n_filters_spat,
+                    filter_time_length=filter_time_length,
+                    bias_time=True,
+                    bias_spat=not self.batch_norm,
+                ),
+            )
+            n_filters_conv = self.n_filters_spat
+        else:
+            self.add_module(
+                "conv_time",
+                nn.Conv2d(
+                    self.n_chans,
+                    self.n_filters_time,
+                    (self.filter_time_length, 1),
+                    stride=(conv_stride, 1),
+                    bias=not self.batch_norm,
+                ),
+            )
+            n_filters_conv = self.n_filters_time
+        if self.batch_norm:
+            self.add_module(
+                "bnorm",
+                nn.BatchNorm2d(
+                    n_filters_conv,
+                    momentum=self.batch_norm_alpha,
+                    affine=True,
+                    eps=1e-5,
+                ),
+            )
+        self.add_module("conv_nonlin", self.first_nonlin())
+        self.add_module(
+            "pool",
+            first_pool_class(
+                kernel_size=(self.pool_time_length, 1), stride=(pool_stride, 1)
+            ),
+        )
+        self.add_module("pool_nonlin", self.first_pool_nonlin())
+
+        def add_conv_pool_block(
+            model, n_filters_before, n_filters, filter_length, block_nr
+        ):
+            suffix = "_{:d}".format(block_nr)
+            self.add_module("drop" + suffix, nn.Dropout(p=self.drop_prob))
+            self.add_module(
+                "conv" + suffix,
+                nn.Conv2d(
+                    n_filters_before,
+                    n_filters,
+                    (filter_length, 1),
+                    stride=(conv_stride, 1),
+                    bias=not self.batch_norm,
+                ),
+            )
+            if self.batch_norm:
+                self.add_module(
+                    "bnorm" + suffix,
+                    nn.BatchNorm2d(
+                        n_filters,
+                        momentum=self.batch_norm_alpha,
+                        affine=True,
+                        eps=1e-5,
+                    ),
+                )
+            self.add_module("nonlin" + suffix, self.later_conv_nonlin())
+
+            self.add_module(
+                "pool" + suffix,
+                later_pool_class(
+                    kernel_size=(self.pool_time_length, 1),
+                    stride=(pool_stride, 1),
+                ),
+            )
+            self.add_module("pool_nonlin" + suffix, self.later_pool_nonlin())
+
+        add_conv_pool_block(
+            self, n_filters_conv, self.n_filters_2, self.filter_length_2, 2
+        )
+        add_conv_pool_block(
+            self, self.n_filters_2, self.n_filters_3, self.filter_length_3, 3
+        )
+        add_conv_pool_block(
+            self, self.n_filters_3, self.n_filters_4, self.filter_length_4, 4
+        )
+
+        # self.add_module('drop_classifier', nn.Dropout(p=self.drop_prob))
+        self.eval()
+        if self.final_conv_length == "auto":
+            self.final_conv_length = self.get_output_shape()[2]
+
+        # Incorporating classification module and subsequent ones in one final layer
+        module = nn.Sequential()
+
+        module.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                self.n_filters_4,
+                self.n_outputs,
+                (self.final_conv_length, 1),
+                bias=True,
+            ),
+        )
+
+        module.add_module("squeeze", SqueezeFinalOutput())
+
+        self.add_module("final_layer", module)
+
+        # Initialization, xavier is same as in our paper...
+        # was default from lasagne
+        init.xavier_uniform_(self.conv_time_spat.conv_time.weight, gain=1)
+        # maybe no bias in case of no split layer and batch norm
+        if self.split_first_layer or (not self.batch_norm):
+            init.constant_(self.conv_time_spat.conv_time.bias, 0)
+        if self.split_first_layer:
+            init.xavier_uniform_(self.conv_time_spat.conv_spat.weight, gain=1)
+            if not self.batch_norm:
+                init.constant_(self.conv_spat.bias, 0)
+        if self.batch_norm:
+            init.constant_(self.bnorm.weight, 1)
+            init.constant_(self.bnorm.bias, 0)
+        param_dict = dict(list(self.named_parameters()))
+        for block_nr in range(2, 5):
+            conv_weight = param_dict["conv_{:d}.weight".format(block_nr)]
+            init.xavier_uniform_(conv_weight, gain=1)
+            if not self.batch_norm:
+                conv_bias = param_dict["conv_{:d}.bias".format(block_nr)]
+                init.constant_(conv_bias, 0)
+            else:
+                bnorm_weight = param_dict["bnorm_{:d}.weight".format(block_nr)]
+                bnorm_bias = param_dict["bnorm_{:d}.bias".format(block_nr)]
+                init.constant_(bnorm_weight, 1)
+                init.constant_(bnorm_bias, 0)
+
+        init.xavier_uniform_(self.final_layer.conv_classifier.weight, gain=1)
+        init.constant_(self.final_layer.conv_classifier.bias, 0)
+
+        self.train()

braindecode/models/deepsleepnet.py

@@ -0,0 +1,295 @@
+# Authors: Théo Gnassounou <theo.gnassounou@inria.fr>
+#
+# License: BSD (3-clause)
+import torch
+import torch.nn as nn
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class DeepSleepNet(EEGModuleMixin, nn.Module):
+    """Sleep staging architecture from Supratak et al. (2017) [Supratak2017]_.
+
+     .. figure:: https://raw.githubusercontent.com/akaraspt/deepsleepnet/refs/heads/master/img/deepsleepnet.png
+        :align: center
+        :alt: DeepSleepNet Architecture
+
+    Convolutional neural network and bidirectional-Long Short-Term
+    for single channels sleep staging described in [Supratak2017]_.
+
+    Parameters
+    ----------
+    activation_large: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+    activation_small: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+    return_feats : bool
+        If True, return the features, i.e. the output of the feature extractor
+        (before the final linear layer). If False, pass the features through
+        the final linear layer.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+
+
+    References
+    ----------
+    .. [Supratak2017] Supratak, A., Dong, H., Wu, C., & Guo, Y. (2017).
+       DeepSleepNet: A model for automatic sleep stage scoring based
+       on raw single-channel EEG. IEEE Transactions on Neural Systems
+       and Rehabilitation Engineering, 25(11), 1998-2008.
+    """
+
+    def __init__(
+        self,
+        n_outputs=5,
+        return_feats=False,
+        n_chans=None,
+        chs_info=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        activation_large: nn.Module = nn.ELU,
+        activation_small: nn.Module = nn.ReLU,
+        drop_prob: float = 0.5,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+        self.cnn1 = _SmallCNN(activation=activation_small, drop_prob=drop_prob)
+        self.cnn2 = _LargeCNN(activation=activation_large, drop_prob=drop_prob)
+        self.dropout = nn.Dropout(0.5)
+        self.bilstm = _BiLSTM(input_size=3072, hidden_size=512, num_layers=2)
+        self.fc = nn.Sequential(
+            nn.Linear(3072, 1024, bias=False), nn.BatchNorm1d(num_features=1024)
+        )
+
+        self.features_extractor = nn.Identity()
+        self.len_last_layer = 1024
+        self.return_feats = return_feats
+
+        # TODO: Add new way to handle return_features == True
+        if not return_feats:
+            self.final_layer = nn.Linear(1024, self.n_outputs)
+        else:
+            self.final_layer = nn.Identity()
+
+    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)
+
+        x1 = self.cnn1(x)
+        x1 = x1.flatten(start_dim=1)
+
+        x2 = self.cnn2(x)
+        x2 = x2.flatten(start_dim=1)
+
+        x = torch.cat((x1, x2), dim=1)
+        x = self.dropout(x)
+        temp = x.clone()
+        temp = self.fc(temp)
+        x = x.unsqueeze(1)
+        x = self.bilstm(x)
+        x = x.squeeze()
+        x = torch.add(x, temp)
+        x = self.dropout(x)
+
+        feats = self.features_extractor(x)
+
+        if self.return_feats:
+            return feats
+        else:
+            return self.final_layer(feats)
+
+
+class _SmallCNN(nn.Module):
+    """
+    Smaller filter sizes to learn temporal information.
+
+    Parameters
+    ----------
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+    drop_prob : float, default=0.5
+        The dropout rate for regularization. Values should be between 0 and 1.
+    """
+
+    def __init__(self, activation: nn.Module = nn.ReLU, drop_prob: float = 0.5):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=64,
+                kernel_size=(1, 50),
+                stride=(1, 6),
+                padding=(0, 22),
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=64),
+            activation(),
+        )
+        self.pool1 = nn.MaxPool2d(kernel_size=(1, 8), stride=(1, 8), padding=(0, 2))
+        self.dropout = nn.Dropout(p=drop_prob)
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=(1, 8),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(1, 8),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(1, 8),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.pool2 = nn.MaxPool2d(kernel_size=(1, 4), stride=(1, 4), padding=(0, 1))
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.dropout(self.pool1(x))
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.pool2(x)
+        return x
+
+
+class _LargeCNN(nn.Module):
+    """
+    Larger filter sizes to learn frequency information.
+
+    Parameters
+    ----------
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    """
+
+    def __init__(self, activation: nn.Module = nn.ELU, drop_prob: float = 0.5):
+        super().__init__()
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=1,
+                out_channels=64,
+                kernel_size=(1, 400),
+                stride=(1, 50),
+                padding=(0, 175),
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=64),
+            activation(),
+        )
+        self.pool1 = nn.MaxPool2d(kernel_size=(1, 4), stride=(1, 4))
+        self.dropout = nn.Dropout(p=drop_prob)
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=(1, 6),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(1, 6),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(1, 6),
+                stride=1,
+                padding="same",
+                bias=False,
+            ),
+            nn.BatchNorm2d(num_features=128),
+            activation(),
+        )
+        self.pool2 = nn.MaxPool2d(kernel_size=(1, 2), stride=(1, 2), padding=(0, 1))
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.dropout(self.pool1(x))
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.pool2(x)
+        return x
+
+
+class _BiLSTM(nn.Module):
+    def __init__(self, input_size, hidden_size, num_layers):
+        super(_BiLSTM, self).__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.lstm = nn.LSTM(
+            input_size,
+            hidden_size,
+            num_layers,
+            batch_first=True,
+            dropout=0.5,
+            bidirectional=True,
+        )
+
+    def forward(self, x):
+        # set initial hidden and cell states
+        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
+        c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
+
+        # forward propagate LSTM
+        out, _ = self.lstm(x, (h0, c0))
+        return out