braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/tidnet.py

@@ -0,0 +1,397 @@
+from math import ceil
+
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+from torch.nn import init
+from torch.nn.utils.parametrizations import weight_norm
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import Ensure4d
+
+
+class TIDNet(EEGModuleMixin, nn.Module):
+    r"""Thinker Invariance DenseNet model from Kostas et al (2020) [TIDNet]_.
+
+    :bdg-success:`Convolution`
+
+    .. figure:: https://content.cld.iop.org/journals/1741-2552/17/5/056008/revision3/jneabb7a7f1_hr.jpg
+        :align: center
+        :alt: TIDNet Architecture
+
+    See [TIDNet]_ for details.
+
+    Parameters
+    ----------
+    s_growth : int
+        DenseNet-style growth factor (added filters per DenseFilter)
+    t_filters : int
+        Number of temporal filters.
+    drop_prob : float
+        Dropout probability
+    pooling : int
+        Max temporal pooling (width and stride)
+    temp_layers : int
+        Number of temporal layers
+    spat_layers : int
+        Number of DenseFilters
+    temp_span : float
+        Percentage of n_times that defines the temporal filter length:
+        temp_len = ceil(temp_span * n_times)
+        e.g A value of 0.05 for temp_span with 1500 n_times will yield a temporal
+        filter of length 75.
+    bottleneck : int
+        Bottleneck factor within Densefilter
+    summary : int
+        Output size of AdaptiveAvgPool1D layer. If set to -1, value will be calculated
+        automatically (n_times // pooling).
+    in_chans :
+        Alias for n_chans.
+    n_classes:
+        Alias for n_outputs.
+    input_window_samples :
+        Alias for n_times.
+    activation: nn.Module, default=nn.LeakyReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.LeakyReLU``.
+
+    Notes
+    -----
+    Code adapted from: https://github.com/SPOClab-ca/ThinkerInvariance/
+
+    References
+    ----------
+    .. [TIDNet] Kostas, D. & Rudzicz, F.
+        Thinker invariance: enabling deep neural networks for BCI across more
+        people.
+        J. Neural Eng. 17, 056008 (2020).
+        doi: 10.1088/1741-2552/abb7a7.
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        input_window_seconds=None,
+        sfreq=None,
+        chs_info=None,
+        s_growth: int = 24,
+        t_filters: int = 32,
+        drop_prob: float = 0.4,
+        pooling: int = 15,
+        temp_layers: int = 2,
+        spat_layers: int = 2,
+        temp_span: float = 0.05,
+        bottleneck: int = 3,
+        summary: int = -1,
+        activation: type[nn.Module] = nn.LeakyReLU,
+    ):
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+            chs_info=chs_info,
+        )
+        del n_outputs, n_chans, n_times, input_window_seconds, sfreq, chs_info
+
+        self.temp_len = ceil(temp_span * self.n_times)
+
+        self.dscnn = _TIDNetFeatures(
+            s_growth=s_growth,
+            t_filters=t_filters,
+            n_chans=self.n_chans,
+            n_times=self.n_times,
+            drop_prob=drop_prob,
+            pooling=pooling,
+            temp_layers=temp_layers,
+            spat_layers=spat_layers,
+            temp_span=temp_span,
+            bottleneck=bottleneck,
+            summary=summary,
+            activation=activation,
+        )
+
+        self._num_features = self.dscnn.num_features
+
+        self.flatten = nn.Flatten(start_dim=1)
+
+        self.final_layer = self._create_classifier(self.num_features, self.n_outputs)
+
+    def _create_classifier(self, incoming: int, n_outputs: int):
+        classifier = nn.Linear(incoming, n_outputs)
+        init.xavier_normal_(classifier.weight)
+        classifier.bias.data.zero_()
+        seq_clf = nn.Sequential(classifier, nn.Identity())
+
+        return seq_clf
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+
+        x = self.dscnn(x)
+        x = self.flatten(x)
+        return self.final_layer(x)
+
+    @property
+    def num_features(self):
+        return self._num_features
+
+
+class _BatchNormZG(nn.BatchNorm2d):
+    def reset_parameters(self):
+        if self.track_running_stats:
+            self.running_mean.zero_()
+            self.running_var.fill_(1)
+        if self.affine:
+            self.weight.data.zero_()
+            self.bias.data.zero_()
+
+
+class _ConvBlock2D(nn.Module):
+    r"""Implements Convolution block with order:
+    Convolution, dropout, activation, batch-norm
+    """
+
+    def __init__(
+        self,
+        in_filters: int,
+        out_filters: int,
+        kernel: tuple[int, int],
+        stride: tuple[int, int] = (1, 1),
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        drop_prob: float = 0.5,
+        batch_norm: bool = True,
+        activation: type[nn.Module] = nn.LeakyReLU,
+        residual: bool = False,
+    ):
+        super().__init__()
+        self.kernel = kernel
+        self.activation = activation()
+        self.residual = residual
+
+        self.conv = nn.Conv2d(
+            in_filters,
+            out_filters,
+            kernel,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=not batch_norm,
+        )
+        self.dropout = nn.Dropout2d(p=float(drop_prob))
+        self.batch_norm = (
+            _BatchNormZG(out_filters)
+            if residual
+            else nn.BatchNorm2d(out_filters)
+            if batch_norm
+            else nn.Identity()
+        )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        res = input
+        input = self.conv(
+            input,
+        )
+        input = self.dropout(input)
+        input = self.activation(input)
+        input = self.batch_norm(input)
+        return input + res if self.residual else input
+
+
+class _DenseFilter(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        growth_rate: int,
+        filter_len: int = 5,
+        drop_prob: float = 0.5,
+        bottleneck: int = 2,
+        activation: type[nn.Module] = nn.LeakyReLU,
+        dim: int = -2,
+    ):
+        super().__init__()
+        dim = dim if dim > 0 else dim + 4
+        if dim < 2 or dim > 3:
+            raise ValueError("Only last two dimensions supported")
+        kernel = (filter_len, 1) if dim == 2 else (1, filter_len)
+
+        self.net = nn.Sequential(
+            nn.BatchNorm2d(in_features),
+            activation(),
+            nn.Conv2d(in_features, bottleneck * growth_rate, 1),
+            nn.BatchNorm2d(bottleneck * growth_rate),
+            activation(),
+            nn.Conv2d(
+                bottleneck * growth_rate,
+                growth_rate,
+                kernel,
+                padding=tuple((k // 2 for k in kernel)),
+            ),
+            nn.Dropout2d(p=float(drop_prob)),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.cat((x, self.net(x)), dim=1)
+
+
+class _DenseSpatialFilter(nn.Module):
+    def __init__(
+        self,
+        n_chans: int,
+        growth: int,
+        depth: int,
+        in_ch: int = 1,
+        bottleneck: int = 4,
+        drop_prob: float = 0.0,
+        activation: type[nn.Module] = nn.LeakyReLU,
+        collapse: bool = True,
+    ):
+        super().__init__()
+        self.net = nn.Sequential(
+            *[
+                _DenseFilter(
+                    in_ch + growth * d,
+                    growth,
+                    bottleneck=bottleneck,
+                    drop_prob=drop_prob,
+                    activation=activation,
+                )
+                for d in range(depth)
+            ]
+        )
+        n_filters = in_ch + growth * depth
+        self.collapse = collapse
+        if collapse:
+            self.channel_collapse = _ConvBlock2D(
+                n_filters, n_filters, (n_chans, 1), drop_prob=0, activation=activation
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if len(x.shape) < 4:
+            x = x.unsqueeze(1).permute([0, 1, 3, 2])
+        x = self.net(x)
+        if self.collapse:
+            return self.channel_collapse(x).squeeze(-2)
+        return x
+
+
+class _TemporalFilter(nn.Module):
+    def __init__(
+        self,
+        n_chans: int,
+        filters: int,
+        depth: int,
+        temp_len: int,
+        drop_prob: float = 0.0,
+        activation: type[nn.Module] = nn.LeakyReLU,
+        residual: str = "netwise",
+    ):
+        super().__init__()
+        temp_len = temp_len + 1 - temp_len % 2
+        self.residual_style = str(residual)
+        net = list()
+
+        for i in range(depth):
+            dil = depth - i
+            conv = weight_norm(
+                nn.Conv2d(
+                    n_chans if i == 0 else filters,
+                    filters,
+                    kernel_size=(1, temp_len),
+                    dilation=dil,
+                    padding=(0, dil * (temp_len - 1) // 2),
+                )
+            )
+            net.append(
+                nn.Sequential(conv, activation(), nn.Dropout2d(p=float(drop_prob)))
+            )
+        if self.residual_style.lower() == "netwise":
+            self.net = nn.Sequential(*net)
+            self.residual = nn.Conv2d(n_chans, filters, (1, 1))
+        elif residual.lower() == "dense":
+            self.net = net
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        style = self.residual_style.lower()
+        if style == "netwise":
+            return self.net(x) + self.residual(x)
+        elif style == "dense":
+            for layer in self.net:
+                x = torch.cat((x, layer(x)), dim=1)
+            return x
+        # TorchScript now knows this path always returns or errors
+        else:
+            # Use an assertion so TorchScript can compile it
+            assert False, f"Unsupported residual style: {self.residual_style}"
+
+
+class _TIDNetFeatures(nn.Module):
+    def __init__(
+        self,
+        s_growth: int,
+        t_filters: int,
+        n_chans: int,
+        n_times: int,
+        drop_prob: float,
+        pooling: int,
+        temp_layers: int,
+        spat_layers: int,
+        temp_span: float,
+        bottleneck: int,
+        summary: int,
+        activation: type[nn.Module] = nn.LeakyReLU,
+    ):
+        super().__init__()
+        self.n_chans = n_chans
+        self.temp_len = ceil(temp_span * n_times)
+
+        self.temporal = nn.Sequential(
+            Ensure4d(),
+            Rearrange("batch C T 1 -> batch 1 C T"),
+            _TemporalFilter(
+                1,
+                t_filters,
+                depth=temp_layers,
+                temp_len=self.temp_len,
+                activation=activation,
+            ),
+            nn.MaxPool2d((1, pooling)),
+            nn.Dropout2d(p=float(drop_prob)),
+        )
+        summary = n_times // pooling if summary == -1 else summary
+
+        self.spatial = _DenseSpatialFilter(
+            n_chans=n_chans,
+            growth=s_growth,
+            depth=spat_layers,
+            in_ch=t_filters,
+            drop_prob=drop_prob,
+            bottleneck=bottleneck,
+            activation=activation,
+        )
+        self.extract_features = nn.Sequential(
+            nn.AdaptiveAvgPool1d(int(summary)), nn.Flatten(start_dim=1)
+        )
+
+        self._num_features = (t_filters + s_growth * spat_layers) * summary
+
+    @property
+    def num_features(self):
+        return self._num_features
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.temporal(x)
+        x = self.spatial(x)
+        return self.extract_features(x)

braindecode/models/tsinception.py

@@ -0,0 +1,295 @@
+# Authors: Bruno Aristimunha <b.aristimunha>
+#
+# License: BSD (3-clause)
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from einops.layers.torch import Rearrange
+from mne.utils import deprecated, warn
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class TSception(EEGModuleMixin, nn.Module):
+    r"""TSception model from Ding et al. (2020) from [ding2020]_.
+
+    :bdg-success:`Convolution`
+
+    TSception: A deep learning framework for emotion detection using EEG.
+
+    .. figure:: https://user-images.githubusercontent.com/58539144/74716976-80415e00-526a-11ea-9433-02ab2b753f6b.PNG
+        :align: center
+        :alt: TSception Architecture
+
+    The model consists of temporal and spatial convolutional layers
+    (Tception and Sception) designed to learn temporal and spatial features
+    from EEG data.
+
+    Parameters
+    ----------
+    number_filter_temp : int
+        Number of temporal convolutional filters.
+    number_filter_spat : int
+        Number of spatial convolutional filters.
+    hidden_size : int
+        Number of units in the hidden fully connected layer.
+    drop_prob : float
+        Dropout rate applied after the hidden layer.
+    activation : nn.Module, optional
+        Activation function class to apply. Should be a PyTorch activation
+        module like ``nn.ReLU`` or ``nn.LeakyReLU``. Default is ``nn.LeakyReLU``.
+    pool_size : int, optional
+        Pooling size for the average pooling layers. Default is 8.
+    inception_windows : list[float], optional
+        List of window sizes (in seconds) for the inception modules.
+        Default is [0.5, 0.25, 0.125].
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors. The modifications are minimal and the model is expected
+    to work as intended. the original code from [code2020]_.
+
+    References
+    ----------
+    .. [ding2020] Ding, Y., Robinson, N., Zeng, Q., Chen, D., Wai, A. A. P.,
+        Lee, T. S., & Guan, C. (2020, July). Tsception: a deep learning framework
+        for emotion detection using EEG. In 2020 international joint conference
+        on neural networks (IJCNN) (pp. 1-7). IEEE.
+    .. [code2020] Ding, Y., Robinson, N., Zeng, Q., Chen, D., Wai, A. A. P.,
+        Lee, T. S., & Guan, C. (2020, July). Tsception: a deep learning framework
+        for emotion detection using EEG.
+        https://github.com/deepBrains/TSception/blob/master/Models.py
+    """
+
+    def __init__(
+        self,
+        # Braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        input_window_seconds=None,
+        chs_info=None,
+        n_times=None,
+        sfreq=None,
+        # Model parameters
+        number_filter_temp: int = 9,
+        number_filter_spat: int = 6,
+        hidden_size: int = 128,
+        drop_prob: float = 0.5,
+        activation: type[nn.Module] = nn.LeakyReLU,
+        pool_size: int = 8,
+        inception_windows: tuple[float, float, float] = (0.5, 0.25, 0.125),
+    ):
+        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.activation = activation
+        self.pool_size = pool_size
+        self.inception_windows = inception_windows
+        self.number_filter_spat = number_filter_spat
+        self.number_filter_temp = number_filter_temp
+        self.drop_prob = drop_prob
+
+        ### Layers
+        self.ensuredim = Rearrange("batch nchans time -> batch 1 nchans time")
+        if self.input_window_seconds < max(self.inception_windows):
+            inception_windows = (
+                self.input_window_seconds,
+                self.input_window_seconds / 2,
+                self.input_window_seconds / 4,
+            )
+            warning_msg = (
+                "Input window size is smaller than the maximum inception window size. "
+                "We are adjusting the input window size to match the maximum inception window size.\n"
+                f"Original input window size: {self.inception_windows}, \n"
+                f"Adjusted inception windows: {inception_windows}"
+            )
+            warn(warning_msg, UserWarning)
+            self.inception_windows = inception_windows
+        # Define temporal convolutional layers (Tception)
+        self.temporal_blocks = nn.ModuleList()
+        for window in self.inception_windows:
+            # 1. Calculate the temporal kernel size for this block
+            kernel_size_t = int(window * self.sfreq)
+
+            # 2. Calculate the output length of the convolution
+            conv_out_len = self.n_times - kernel_size_t + 1
+
+            # 3. Ensure the pooling size is not larger than the conv output
+            #    and is at least 1.
+            dynamic_pool_size = max(1, min(self.pool_size, conv_out_len))
+
+            # 4. Create the block with the dynamic pooling size
+            block = self._conv_block(
+                in_channels=1,
+                out_channels=self.number_filter_temp,
+                kernel_size=(1, kernel_size_t),
+                stride=1,
+                pool_size=dynamic_pool_size,  # Use the dynamic size
+                activation=self.activation,
+            )
+            self.temporal_blocks.append(block)
+
+        self.batch_temporal_lay = nn.BatchNorm2d(self.number_filter_temp)
+
+        # Define spatial convolutional layers (Sception)
+
+        pool_size_spat = self.pool_size // 4
+
+        self.spatial_block_1 = self._conv_block(
+            in_channels=self.number_filter_temp,
+            out_channels=self.number_filter_spat,
+            kernel_size=(self.n_chans, 1),
+            stride=1,
+            pool_size=pool_size_spat,
+            activation=self.activation,
+        )
+
+        kernel_size_spat_2 = (max(1, self.n_chans // 2), 1)
+
+        self.spatial_block_2 = self._conv_block(
+            in_channels=self.number_filter_temp,
+            out_channels=self.number_filter_spat,
+            kernel_size=kernel_size_spat_2,
+            stride=kernel_size_spat_2,
+            pool_size=pool_size_spat,
+            activation=self.activation,
+        )
+        self.batch_spatial_lay = nn.BatchNorm2d(self.number_filter_spat)
+
+        # Calculate the size of the features after convolution and pooling layers
+        self.feature_size = self._calculate_feature_size()
+        # self.feature_size = self.number_filter_spat *
+        # Define the final classification layers
+
+        self.dense_layer = nn.Sequential(
+            nn.Linear(self.feature_size, hidden_size),
+            self.activation(),
+            nn.Dropout(self.drop_prob),
+        )
+
+        self.final_layer = nn.Linear(hidden_size, self.n_outputs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the TSception model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, n_channels, n_times).
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (batch_size, n_classes).
+        """
+        # Temporal Convolution
+        # shape: (batch_size, n_channels, n_times)
+        x = self.ensuredim(x)
+        # shape: (batch_size, 1, n_channels, n_times)
+
+        t_features = [layer(x) for layer in self.temporal_blocks]
+        # shape: (batch_size, number_filter_temp, n_channels,
+        #
+        t_out = torch.cat(t_features, dim=-1)
+
+        t_out = self.batch_temporal_lay(t_out)
+
+        # Spatial Convolution
+        s_out1 = self.spatial_block_1(t_out)
+        s_out2 = self.spatial_block_2(t_out)
+        s_out = torch.cat((s_out1, s_out2), dim=2)
+        s_out = self.batch_spatial_lay(s_out)
+
+        # Flatten and apply final layers
+        s_out = s_out.view(s_out.size(0), -1)
+        output = self.dense_layer(s_out)
+        output = self.final_layer(output)
+        return output
+
+    def _calculate_feature_size(self) -> int:
+        """
+        Calculates the size of the features after convolution and pooling layers.
+
+        Returns
+        -------
+        int
+            Flattened size of the features after convolution and pooling layers.
+        """
+        with torch.no_grad():
+            dummy_input = torch.ones(1, 1, self.n_chans, self.n_times)
+            t_features = [layer(dummy_input) for layer in self.temporal_blocks]
+            t_out = torch.cat(t_features, dim=-1)
+            t_out = self.batch_temporal_lay(t_out)
+
+            s_out1 = self.spatial_block_1(t_out)
+            s_out2 = self.spatial_block_2(t_out)
+            s_out = torch.cat((s_out1, s_out2), dim=2)
+            s_out = self.batch_spatial_lay(s_out)
+
+            feature_size = s_out.view(1, -1).size(1)
+        return feature_size
+
+    @staticmethod
+    def _conv_block(
+        in_channels: int,
+        out_channels: int,
+        kernel_size: tuple,
+        stride: tuple[int, int] | int,
+        pool_size: int,
+        activation: nn.Module,
+    ) -> nn.Sequential:
+        """
+        Creates a convolutional block with Conv2d, activation, and AvgPool2d layers.
+
+        Parameters
+        ----------
+        in_channels : int
+            Number of input channels.
+        out_channels : int
+            Number of output channels.
+        kernel_size : tuple
+            Size of the convolutional kernel.
+        stride : int
+            Stride of the convolution.
+        pool_size : int
+            Size of the pooling kernel.
+        activation : nn.Module
+            Activation function class.
+
+        Returns
+        -------
+        nn.Sequential
+            A sequential container of the convolutional block.
+        """
+        return nn.Sequential(
+            nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=0,
+            ),
+            activation(),
+            nn.AvgPool2d(kernel_size=(1, pool_size), stride=(1, pool_size)),
+        )
+
+
+@deprecated(
+    "`TSceptionV1` was renamed to `TSception` in v1.12; "
+    "this alias will be removed in v1.14."
+)
+class TSceptionV1(TSception):
+    r"""Deprecated alias for TSception."""
+
+    pass