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

braindecode/models/atcnet.py

@@ -0,0 +1,652 @@
+# Authors: Cedric Rommel <cedric.rommel@inria.fr>
+#
+# License: BSD (3-clause)
+import math
+
+import torch
+from einops.layers.torch import Rearrange
+from torch import nn
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import CausalConv1d, Ensure4d, MaxNormLinear
+
+
+class ATCNet(EEGModuleMixin, nn.Module):
+    """ATCNet model from Altaheri et al. (2022) [1]_
+
+    Pytorch implementation based on official tensorflow code [2]_.
+
+    .. figure:: https://user-images.githubusercontent.com/25565236/185449791-e8539453-d4fa-41e1-865a-2cf7e91f60ef.png
+       :align: center
+       :alt: ATCNet Architecture
+
+    Parameters
+    ----------
+    input_window_seconds : float, optional
+        Time length of inputs, in seconds. Defaults to 4.5 s, as in BCI-IV 2a
+        dataset.
+    sfreq : int, optional
+        Sampling frequency of the inputs, in Hz. Default to 250 Hz, as in
+        BCI-IV 2a dataset.
+    conv_block_n_filters : int
+        Number temporal filters in the first convolutional layer of the
+        convolutional block, denoted F1 in figure 2 of the paper [1]_. Defaults
+        to 16 as in [1]_.
+    conv_block_kernel_length_1 : int
+        Length of temporal filters in the first convolutional layer of the
+        convolutional block, denoted Kc in table 1 of the paper [1]_. Defaults
+        to 64 as in [1]_.
+    conv_block_kernel_length_2 : int
+        Length of temporal filters in the last convolutional layer of the
+        convolutional block. Defaults to 16 as in [1]_.
+    conv_block_pool_size_1 : int
+        Length of first average pooling kernel in the convolutional block.
+        Defaults to 8 as in [1]_.
+    conv_block_pool_size_2 : int
+        Length of first average pooling kernel in the convolutional block,
+        denoted P2 in table 1 of the paper [1]_. Defaults to 7 as in [1]_.
+    conv_block_depth_mult : int
+        Depth multiplier of depthwise convolution in the convolutional block,
+        denoted D in table 1 of the paper [1]_. Defaults to 2 as in [1]_.
+    conv_block_dropout : float
+        Dropout probability used in the convolution block, denoted pc in
+        table 1 of the paper [1]_. Defaults to 0.3 as in [1]_.
+    n_windows : int
+        Number of sliding windows, denoted n in [1]_. Defaults to 5 as in [1]_.
+    att_head_dim : int
+        Embedding dimension used in each self-attention head, denoted dh in
+        table 1 of the paper [1]_. Defaults to 8 as in [1]_.
+    att_num_heads : int
+        Number of attention heads, denoted H in table 1 of the paper [1]_.
+        Defaults to 2 as in [1]_.
+    att_dropout : float
+        Dropout probability used in the attention block, denoted pa in table 1
+        of the paper [1]_. Defaults to 0.5 as in [1]_.
+    tcn_depth : int
+        Depth of Temporal Convolutional Network block (i.e. number of TCN
+        Residual blocks), denoted L in table 1 of the paper [1]_. Defaults to 2
+        as in [1]_.
+    tcn_kernel_size : int
+        Temporal kernel size used in TCN block, denoted Kt in table 1 of the
+        paper [1]_. Defaults to 4 as in [1]_.
+    tcn_n_filters : int
+        Number of filters used in TCN convolutional layers (Ft). Defaults to
+        32 as in [1]_.
+    tcn_dropout : float
+        Dropout probability used in the TCN block, denoted pt in table 1
+        of the paper [1]_. Defaults to 0.3 as in [1]_.
+    tcn_activation : torch.nn.Module
+        Nonlinear activation to use. Defaults to nn.ELU().
+    concat : bool
+        When ``True``, concatenates each slidding window embedding before
+        feeding it to a fully-connected layer, as done in [1]_. When ``False``,
+        maps each slidding window to `n_outputs` logits and average them.
+        Defaults to ``False`` contrary to what is reported in [1]_, but
+        matching what the official code does [2]_.
+    max_norm_const : float
+        Maximum L2-norm constraint imposed on weights of the last
+        fully-connected layer. Defaults to 0.25.
+
+
+    References
+    ----------
+    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman,
+        Physics-informed attention temporal convolutional network for EEG-based
+        motor imagery classification in IEEE Transactions on Industrial Informatics,
+        2022, doi: 10.1109/TII.2022.3197419.
+    .. [2] EEE-ATCNet implementation.
+       https://github.com/Altaheri/EEG-ATCNet/blob/main/models.py
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        input_window_seconds=None,
+        sfreq=250.0,
+        conv_block_n_filters=16,
+        conv_block_kernel_length_1=64,
+        conv_block_kernel_length_2=16,
+        conv_block_pool_size_1=8,
+        conv_block_pool_size_2=7,
+        conv_block_depth_mult=2,
+        conv_block_dropout=0.3,
+        n_windows=5,
+        att_head_dim=8,
+        att_num_heads=2,
+        att_drop_prob=0.5,
+        tcn_depth=2,
+        tcn_kernel_size=4,
+        tcn_n_filters=32,
+        tcn_drop_prob=0.3,
+        tcn_activation: nn.Module = nn.ELU,
+        concat=False,
+        max_norm_const=0.25,
+        chs_info=None,
+        n_times=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
+        self.conv_block_n_filters = conv_block_n_filters
+        self.conv_block_kernel_length_1 = conv_block_kernel_length_1
+        self.conv_block_kernel_length_2 = conv_block_kernel_length_2
+        self.conv_block_pool_size_1 = conv_block_pool_size_1
+        self.conv_block_pool_size_2 = conv_block_pool_size_2
+        self.conv_block_depth_mult = conv_block_depth_mult
+        self.conv_block_dropout = conv_block_dropout
+        self.n_windows = n_windows
+        self.att_head_dim = att_head_dim
+        self.att_num_heads = att_num_heads
+        self.att_dropout = att_drop_prob
+        self.tcn_depth = tcn_depth
+        self.tcn_kernel_size = tcn_kernel_size
+        self.tcn_n_filters = tcn_n_filters
+        self.tcn_dropout = tcn_drop_prob
+        self.tcn_activation = tcn_activation
+        self.concat = concat
+        self.max_norm_const = max_norm_const
+
+        map = dict()
+        for w in range(self.n_windows):
+            map[f"max_norm_linears.[{w}].weight"] = f"final_layer.[{w}].weight"
+            map[f"max_norm_linears.[{w}].bias"] = f"final_layer.[{w}].bias"
+        self.mapping = map
+
+        # Check later if we want to keep the Ensure4d. Not sure if we can
+        # remove it or replace it with eipsum.
+        self.ensuredims = Ensure4d()
+        self.dimshuffle = Rearrange("batch C T 1 -> batch 1 T C")
+
+        self.conv_block = _ConvBlock(
+            n_channels=self.n_chans,  # input shape: (batch_size, 1, T, C)
+            n_filters=conv_block_n_filters,
+            kernel_length_1=conv_block_kernel_length_1,
+            kernel_length_2=conv_block_kernel_length_2,
+            pool_size_1=conv_block_pool_size_1,
+            pool_size_2=conv_block_pool_size_2,
+            depth_mult=conv_block_depth_mult,
+            dropout=conv_block_dropout,
+        )
+
+        self.F2 = int(conv_block_depth_mult * conv_block_n_filters)
+        self.Tc = int(self.n_times / (conv_block_pool_size_1 * conv_block_pool_size_2))
+        self.Tw = self.Tc - self.n_windows + 1
+
+        self.attention_blocks = nn.ModuleList(
+            [
+                _AttentionBlock(
+                    in_shape=self.F2,
+                    head_dim=self.att_head_dim,
+                    num_heads=att_num_heads,
+                    dropout=att_drop_prob,
+                )
+                for _ in range(self.n_windows)
+            ]
+        )
+
+        self.temporal_conv_nets = nn.ModuleList(
+            [
+                nn.Sequential(
+                    *[
+                        _TCNResidualBlock(
+                            in_channels=self.F2,
+                            kernel_size=tcn_kernel_size,
+                            n_filters=tcn_n_filters,
+                            dropout=tcn_drop_prob,
+                            activation=tcn_activation,
+                            dilation=2**i,
+                        )
+                        for i in range(tcn_depth)
+                    ]
+                )
+                for _ in range(self.n_windows)
+            ]
+        )
+
+        if self.concat:
+            self.final_layer = nn.ModuleList(
+                [
+                    MaxNormLinear(
+                        in_features=self.F2 * self.n_windows,
+                        out_features=self.n_outputs,
+                        max_norm_val=self.max_norm_const,
+                    )
+                ]
+            )
+        else:
+            self.final_layer = nn.ModuleList(
+                [
+                    MaxNormLinear(
+                        in_features=self.F2,
+                        out_features=self.n_outputs,
+                        max_norm_val=self.max_norm_const,
+                    )
+                    for _ in range(self.n_windows)
+                ]
+            )
+
+        self.out_fun = nn.Identity()
+
+    def forward(self, X):
+        # Dimension: (batch_size, C, T)
+        X = self.ensuredims(X)
+        # Dimension: (batch_size, C, T, 1)
+        X = self.dimshuffle(X)
+        # Dimension: (batch_size, 1, T, C)
+
+        # ----- Sliding window -----
+        conv_feat = self.conv_block(X)
+        # Dimension: (batch_size, F2, Tc, 1)
+        conv_feat = conv_feat.view(-1, self.F2, self.Tc)
+        # Dimension: (batch_size, F2, Tc)
+
+        # ----- Sliding window -----
+        sw_concat: list[torch.Tensor] = []  # to store sliding window outputs
+        # for w in range(self.n_windows):
+        for idx, (attention, tcn_module, final_layer) in enumerate(
+            zip(self.attention_blocks, self.temporal_conv_nets, self.final_layer)
+        ):
+            conv_feat_w = conv_feat[..., idx : idx + self.Tw]
+            # Dimension: (batch_size, F2, Tw)
+
+            # ----- Attention block -----
+            att_feat = attention(conv_feat_w)
+            # Dimension: (batch_size, F2, Tw)
+
+            # ----- Temporal convolutional network (TCN) -----
+            tcn_feat = tcn_module(att_feat)[..., -1]
+            # Dimension: (batch_size, F2)
+
+            # Outputs of sliding window can be either averaged after being
+            # mapped by dense layer or concatenated then mapped by a dense
+            # layer
+            if not self.concat:
+                tcn_feat = final_layer(tcn_feat)
+
+            sw_concat.append(tcn_feat)
+
+        # ----- Aggregation and prediction -----
+        if self.concat:
+            sw_concat_agg = torch.cat(sw_concat, dim=1)
+            sw_concat_agg = self.final_layer[0](sw_concat_agg)
+        else:
+            if len(sw_concat) > 1:  # more than one window
+                sw_concat_agg = torch.stack(sw_concat, dim=0)
+                sw_concat_agg = torch.mean(sw_concat_agg, dim=0)
+            else:  # one window (# windows = 1)
+                sw_concat_agg = sw_concat[0]
+
+        return self.out_fun(sw_concat_agg)
+
+
+class _ConvBlock(nn.Module):
+    """Convolutional block proposed in ATCNet [1]_, inspired by the EEGNet
+    architecture [2]_.
+
+    References
+    ----------
+    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman, "Physics-informed
+           attention temporal convolutional network for EEG-based motor imagery
+           classification," in IEEE Transactions on Industrial Informatics,
+           2022, doi: 10.1109/TII.2022.3197419.
+    .. [2] Lawhern, V. J., Solon, A. J., Waytowich, N. R., Gordon,
+           S. M., Hung, C. P., & Lance, B. J. (2018).
+           EEGNet: A Compact Convolutional Network for EEG-based
+           Brain-Computer Interfaces.
+           arXiv preprint arXiv:1611.08024.
+    """
+
+    def __init__(
+        self,
+        n_channels,
+        n_filters=16,
+        kernel_length_1=64,
+        kernel_length_2=16,
+        pool_size_1=8,
+        pool_size_2=7,
+        depth_mult=2,
+        dropout=0.3,
+    ):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(
+            in_channels=1,
+            out_channels=n_filters,
+            kernel_size=(kernel_length_1, 1),
+            padding="same",
+            bias=False,
+        )
+
+        self.bn1 = nn.BatchNorm2d(num_features=n_filters, eps=1e-4)
+
+        n_depth_kernels = n_filters * depth_mult
+        self.conv2 = nn.Conv2d(
+            in_channels=n_filters,
+            out_channels=n_depth_kernels,
+            groups=n_filters,
+            kernel_size=(1, n_channels),
+            padding="valid",
+            bias=False,
+        )
+
+        self.bn2 = nn.BatchNorm2d(num_features=n_depth_kernels, eps=1e-4)
+
+        self.activation2 = nn.ELU()
+
+        self.pool2 = nn.AvgPool2d(kernel_size=(pool_size_1, 1))
+
+        self.drop2 = nn.Dropout2d(dropout)
+
+        self.conv3 = nn.Conv2d(
+            in_channels=n_depth_kernels,
+            out_channels=n_depth_kernels,
+            kernel_size=(kernel_length_2, 1),
+            padding="same",
+            bias=False,
+        )
+
+        self.bn3 = nn.BatchNorm2d(num_features=n_depth_kernels, eps=1e-4)
+
+        self.activation3 = nn.ELU()
+
+        self.pool3 = nn.AvgPool2d(kernel_size=(pool_size_2, 1))
+
+        self.drop3 = nn.Dropout2d(dropout)
+
+    def forward(self, X):
+        # ----- Temporal convolution -----
+        # Dimension: (batch_size, 1, T, C)
+        X = self.conv1(X)
+        X = self.bn1(X)
+        # Dimension: (batch_size, F1, T, C)
+
+        # ----- Depthwise channels convolution -----
+        X = self.conv2(X)
+        X = self.bn2(X)
+        X = self.activation2(X)
+        # Dimension: (batch_size, F1*D, T, 1)
+        X = self.pool2(X)
+        X = self.drop2(X)
+        # Dimension: (batch_size, F1*D, T/P1, 1)
+
+        # ----- "Spatial" convolution -----
+        X = self.conv3(X)
+        X = self.bn3(X)
+        X = self.activation3(X)
+        # Dimension: (batch_size, F1*D, T/P1, 1)
+        X = self.pool3(X)
+        X = self.drop3(X)
+        # Dimension: (batch_size, F1*D, T/(P1*P2), 1)
+
+        return X
+
+
+class _AttentionBlock(nn.Module):
+    """Multi Head self Attention (MHA) block used in ATCNet [1]_, inspired from
+    [2]_.
+
+    References
+    ----------
+    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman, "Physics-informed
+           attention temporal convolutional network for EEG-based motor imagery
+           classification," in IEEE Transactions on Industrial Informatics,
+           2022, doi: 10.1109/TII.2022.3197419.
+    .. [2] Vaswani, A. et al., "Attention is all you need",
+           in Advances in neural information processing systems, 2017.
+    """
+
+    def __init__(
+        self,
+        in_shape=32,
+        head_dim=8,
+        num_heads=2,
+        dropout=0.5,
+    ):
+        super().__init__()
+        self.in_shape = in_shape
+        self.head_dim = head_dim
+        self.num_heads = num_heads
+
+        # Puts time dimension at -2 and feature dim at -1
+        self.dimshuffle = Rearrange("batch C T -> batch T C")
+
+        # Layer normalization
+        self.ln = nn.LayerNorm(normalized_shape=in_shape, eps=1e-6)
+
+        # Multi-head self-attention layer
+        # (We had to reimplement it since the original code is in tensorflow,
+        # where it is possible to have an embedding dimension different than
+        # the input and output dimensions, which is not possible in pytorch.)
+        self.mha = _MHA(
+            input_dim=in_shape,
+            head_dim=head_dim,
+            output_dim=in_shape,
+            num_heads=num_heads,
+            dropout=dropout,
+        )
+
+        # XXX: This line in the official code is weird, as there is already
+        # dropout in the MultiheadAttention layer. They also don't mention
+        # any additional dropout between the attention block and TCN in the
+        # paper. We are adding it here however to follo so we are removing this
+        # for now.
+        self.drop = nn.Dropout(0.3)
+
+    def forward(self, X):
+        # Dimension: (batch_size, F2, Tw)
+        X = self.dimshuffle(X)
+        # Dimension: (batch_size, Tw, F2)
+
+        # ----- Layer norm -----
+        out = self.ln(X)
+
+        # ----- Self-Attention -----
+        out = self.mha(out, out, out)
+        # Dimension: (batch_size, Tw, F2)
+
+        # XXX In the paper fig. 1, it is drawn that layer normalization is
+        # performed before the skip connection, while it is done afterwards
+        # in the official code. Here we follow the code.
+
+        # ----- Skip connection -----
+        out = X + self.drop(out)
+
+        # Move back to shape (batch_size, F2, Tw) from the beginning
+        return self.dimshuffle(out)
+
+
+class _TCNResidualBlock(nn.Module):
+    """Modified TCN Residual block as proposed in [1]_. Inspired from
+    Temporal Convolutional Networks (TCN) [2]_.
+
+    References
+    ----------
+    .. [1] H. Altaheri, G. Muhammad and M. Alsulaiman, "Physics-informed
+           attention temporal convolutional network for EEG-based motor imagery
+           classification," in IEEE Transactions on Industrial Informatics,
+           2022, doi: 10.1109/TII.2022.3197419.
+    .. [2] Bai, S., Kolter, J. Z., & Koltun, V.
+           "An empirical evaluation of generic convolutional and recurrent
+           networks for sequence modeling", 2018.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        kernel_size=4,
+        n_filters=32,
+        dropout=0.3,
+        activation: nn.Module = nn.ELU,
+        dilation=1,
+    ):
+        super().__init__()
+        self.activation = activation()
+        self.dilation = dilation
+        self.dropout = dropout
+        self.n_filters = n_filters
+        self.kernel_size = kernel_size
+        self.in_channels = in_channels
+
+        self.conv1 = CausalConv1d(
+            in_channels=in_channels,
+            out_channels=n_filters,
+            kernel_size=kernel_size,
+            dilation=dilation,
+        )
+        nn.init.kaiming_uniform_(self.conv1.weight)
+
+        self.bn1 = nn.BatchNorm1d(n_filters)
+
+        self.drop1 = nn.Dropout(dropout)
+
+        self.conv2 = CausalConv1d(
+            in_channels=n_filters,
+            out_channels=n_filters,
+            kernel_size=kernel_size,
+            dilation=dilation,
+        )
+        nn.init.kaiming_uniform_(self.conv2.weight)
+
+        self.bn2 = nn.BatchNorm1d(n_filters)
+
+        self.drop2 = nn.Dropout(dropout)
+
+        # Reshape the input for the residual connection when necessary
+        if in_channels != n_filters:
+            self.reshaping_conv = nn.Conv1d(
+                n_filters,
+                kernel_size=1,
+                padding="same",
+            )
+        else:
+            self.reshaping_conv = nn.Identity()
+
+    def forward(self, X):
+        # Dimension: (batch_size, F2, Tw)
+        # ----- Double dilated convolutions -----
+        out = self.conv1(X)
+        out = self.bn1(out)
+        out = self.activation(out)
+        out = self.drop1(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.activation(out)
+        out = self.drop2(out)
+
+        out = self.reshaping_conv(out)
+
+        # ----- Residual connection -----
+        out = X + out
+
+        return self.activation(out)
+
+
+class _MHA(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        head_dim: int,
+        output_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+    ):
+        """Multi-head Attention
+
+        The difference between this module and torch.nn.MultiheadAttention is
+        that this module supports embedding dimensions different then input
+        and output ones. It also does not support sequences of different
+        length.
+
+        Parameters
+        ----------
+        input_dim : int
+            Dimension of query, key and value inputs.
+        head_dim : int
+            Dimension of embed query, key and value in each head,
+            before computing attention.
+        output_dim : int
+            Output dimension.
+        num_heads : int
+            Number of heads in the multi-head architecture.
+        dropout : float, optional
+            Dropout probability on output weights. Default: 0.0 (no dropout).
+        """
+
+        super(_MHA, self).__init__()
+
+        self.input_dim = input_dim
+        self.head_dim = head_dim
+        # typical choice for the split dimension of the heads
+        self.embed_dim = head_dim * num_heads
+
+        # embeddings for multi-head projections
+        self.fc_q = nn.Linear(input_dim, self.embed_dim)
+        self.fc_k = nn.Linear(input_dim, self.embed_dim)
+        self.fc_v = nn.Linear(input_dim, self.embed_dim)
+
+        # output mapping
+        self.fc_o = nn.Linear(self.embed_dim, output_dim)
+
+        # dropout
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor
+    ) -> torch.Tensor:
+        """Compute MHA(Q, K, V)
+
+        Parameters
+        ----------
+        Q: torch.Tensor of size (batch_size, seq_len, input_dim)
+            Input query (Q) sequence.
+        K: torch.Tensor of size (batch_size, seq_len, input_dim)
+            Input key (K) sequence.
+        V: torch.Tensor of size (batch_size, seq_len, input_dim)
+            Input value (V) sequence.
+
+        Returns
+        -------
+        O: torch.Tensor of size (batch_size, seq_len, output_dim)
+            Output MHA(Q, K, V)
+        """
+        assert Q.shape[-1] == K.shape[-1] == V.shape[-1] == self.input_dim
+
+        batch_size, _, _ = Q.shape
+
+        # embedding for multi-head projections (masked or not)
+        Q = self.fc_q(Q)  # (B, S, D)
+        K, V = self.fc_k(K), self.fc_v(V)  # (B, S, D)
+
+        # Split into num_head vectors (num_heads * batch_size, n/m, head_dim)
+        Q_ = torch.cat(Q.split(self.head_dim, -1), 0)  # (B', S, D')
+        K_ = torch.cat(K.split(self.head_dim, -1), 0)  # (B', S, D')
+        V_ = torch.cat(V.split(self.head_dim, -1), 0)  # (B', S, D')
+
+        # Attention weights of size (num_heads * batch_size, n, m):
+        # measures how similar each pair of Q and K is.
+        W = torch.softmax(
+            Q_.bmm(K_.transpose(-2, -1)) / math.sqrt(self.head_dim),
+            -1,  # (B', D', S)
+        )  # (B', N, M)
+
+        # Multihead output (batch_size, seq_len, dim):
+        # weighted sum of V where a value gets more weight if its corresponding
+        # key has larger dot product with the query.
+        H = torch.cat(
+            (W.bmm(V_)).split(  # (B', S, S)  # (B', S, D')
+                batch_size, 0
+            ),  # [(B, S, D')] * num_heads
+            -1,
+        )  # (B, S, D)
+
+        out = self.fc_o(H)
+
+        return self.dropout(out)