# Authors: Tao Yang <sheeptao@outlook.com>
# Bruno Aristimunha <b.aristimunha@gmail.com> (braindecode adaptation)
#
import torch
import torch.nn as nn
from einops.layers.torch import Rearrange
from typing import List, Type, Union, Tuple, Optional, Dict
from braindecode.models.base import EEGModuleMixin
class _TSConv(nn.Sequential):
"""
Time-Distributed Separable Convolution block.
The architecture consists of:
- **Temporal Convolution**
- **Batch Normalization**
- **Depthwise Spatial Convolution**
- **Batch Normalization**
- **Activation Function**
- **First Pooling Layer**
- **Dropout**
- **Depthwise Temporal Convolution**
- **Batch Normalization**
- **Activation Function**
- **Second Pooling Layer**
- **Dropout**
Parameters
----------
n_channels : int
Number of input channels (EEG channels).
n_filters : int
Number of filters for the convolution layers.
conv1_kernel_size : int
Kernel size for the first convolution layer.
conv2_kernel_size : int
Kernel size for the second convolution layer.
depth_multiplier : int
Depth multiplier for depthwise convolution.
pool1_size : int
Kernel size for the first pooling layer.
pool2_size : int
Kernel size for the second pooling layer.
drop_prob : float
Dropout probability.
activation : Type[nn.Module], optional
Activation function class to use, by default nn.ELU.
"""
def __init__(
self,
n_channels: int,
n_filters: int,
conv1_kernel_size: int,
conv2_kernel_size: int,
depth_multiplier: int,
pool1_size: int,
pool2_size: int,
drop_prob: float,
activation: Type[nn.Module] = nn.ELU,
):
super().__init__(
nn.Conv2d(
in_channels=1,
out_channels=n_filters,
kernel_size=(1, conv1_kernel_size),
padding="same",
bias=False,
),
nn.BatchNorm2d(n_filters),
nn.Conv2d(
in_channels=n_filters,
out_channels=n_filters * depth_multiplier,
kernel_size=(n_channels, 1),
groups=n_filters,
bias=False,
),
nn.BatchNorm2d(n_filters * depth_multiplier),
activation(),
nn.AvgPool2d(kernel_size=(1, pool1_size)),
nn.Dropout(drop_prob),
nn.Conv2d(
in_channels=n_filters * depth_multiplier,
out_channels=n_filters * depth_multiplier,
kernel_size=(1, conv2_kernel_size),
padding="same",
groups=n_filters * depth_multiplier,
bias=False,
),
nn.BatchNorm2d(n_filters * depth_multiplier),
activation(),
nn.AvgPool2d(kernel_size=(1, pool2_size)),
nn.Dropout(drop_prob),
)
class _PositionalEncoding(nn.Module):
"""
Positional encoding module that adds learnable positional embeddings.
Parameters
----------
seq_length : int
Sequence length.
d_model : int
Dimensionality of the model.
"""
def __init__(self, seq_length: int, d_model: int) -> None:
super().__init__()
self.seq_length = seq_length
self.d_model = d_model
self.pe = nn.Parameter(torch.zeros(1, seq_length, d_model))
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.pe
return x
class _Transformer(nn.Module):
"""
Transformer encoder module with learnable class token and positional encoding.
Parameters
----------
seq_length : int
Sequence length of the input.
d_model : int
Dimensionality of the model.
num_heads : int
Number of heads in the multihead attention.
feedforward_ratio : float
Ratio to compute the dimension of the feedforward network.
drop_prob : float, optional
Dropout probability, by default 0.5.
num_layers : int, optional
Number of transformer encoder layers, by default 4.
"""
def __init__(
self,
seq_length: int,
d_model: int,
num_heads: int,
feedforward_ratio: float,
drop_prob: float = 0.5,
num_layers: int = 4,
) -> None:
super().__init__()
self.cls_embedding = nn.Parameter(torch.zeros(1, 1, d_model))
self.pos_embedding = _PositionalEncoding(seq_length + 1, d_model)
dim_ff = int(d_model * feedforward_ratio)
self.dropout = nn.Dropout(drop_prob)
self.trans = nn.TransformerEncoder(
nn.TransformerEncoderLayer(
d_model,
num_heads,
dim_ff,
drop_prob,
batch_first=True,
norm_first=True,
),
num_layers,
norm=nn.LayerNorm(d_model),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
batch_size = x.shape[0]
x = torch.cat((self.cls_embedding.expand(batch_size, -1, -1), x), dim=1)
x = self.pos_embedding(x)
x = self.dropout(x)
return self.trans(x)[:, 0]
class _DenseLayers(nn.Sequential):
"""
Final classification layers.
Parameters
----------
linear_in : int
Input dimension to the linear layer.
n_classes : int
Number of output classes.
"""
def __init__(self, linear_in: int, n_classes: int):
super().__init__(
nn.Flatten(),
nn.Linear(linear_in, n_classes),
)
class MSVTNet(EEGModuleMixin, nn.Module):
"""MSVTNet model from Liu K et al (2024) from [msvt2024]_.
This model implements a multi-scale convolutional transformer network
for EEG signal classification, as described in [msvt2024]_.
.. figure:: https://raw.githubusercontent.com/SheepTAO/MSVTNet/refs/heads/main/MSVTNet_Arch.png
:align: center
:alt: MSVTNet Architecture
Parameters
----------
n_filters_list : List[int], optional
List of filter numbers for each TSConv block, by default (9, 9, 9, 9).
conv1_kernels_size : List[int], optional
List of kernel sizes for the first convolution in each TSConv block,
by default (15, 31, 63, 125).
conv2_kernel_size : int, optional
Kernel size for the second convolution in TSConv blocks, by default 15.
depth_multiplier : int, optional
Depth multiplier for depthwise convolution, by default 2.
pool1_size : int, optional
Pooling size for the first pooling layer in TSConv blocks, by default 8.
pool2_size : int, optional
Pooling size for the second pooling layer in TSConv blocks, by default 7.
drop_prob : float, optional
Dropout probability for convolutional layers, by default 0.3.
num_heads : int, optional
Number of attention heads in the transformer encoder, by default 8.
feedforward_ratio : float, optional
Ratio to compute feedforward dimension in the transformer, by default 1.
drop_prob_trans : float, optional
Dropout probability for the transformer, by default 0.5.
num_layers : int, optional
Number of transformer encoder layers, by default 2.
activation : Type[nn.Module], optional
Activation function class to use, by default nn.ELU.
return_features : bool, optional
Whether to return predictions from branch classifiers, by default False.
Notes
-----
This implementation is not guaranteed to be correct, has not been checked
by original authors, only reimplemented based on the original code [msvt2024code]_.
References
----------
.. [msvt2024] Liu, K., et al. (2024). MSVTNet: Multi-Scale Vision
Transformer Neural Network for EEG-Based Motor Imagery Decoding.
IEEE Journal of Biomedical an Health Informatics.
.. [msvt2024code] Liu, K., et al. (2024). MSVTNet: Multi-Scale Vision
Transformer Neural Network for EEG-Based Motor Imagery Decoding.
Source Code: https://github.com/SheepTAO/MSVTNet
"""
def __init__(
self,
# braindecode parameters
n_chans: Optional[int] = None,
n_outputs: Optional[int] = None,
n_times: Optional[int] = None,
input_window_seconds: Optional[float] = None,
sfreq: Optional[float] = None,
chs_info: Optional[List[Dict]] = None,
# Model's parameters
n_filters_list: Tuple[int, ...] = (9, 9, 9, 9),
conv1_kernels_size: Tuple[int, ...] = (15, 31, 63, 125),
conv2_kernel_size: int = 15,
depth_multiplier: int = 2,
pool1_size: int = 8,
pool2_size: int = 7,
drop_prob: float = 0.3,
num_heads: int = 8,
feedforward_ratio: float = 1,
drop_prob_trans: float = 0.5,
num_layers: int = 2,
activation: Type[nn.Module] = nn.ELU,
return_features: bool = False,
):
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.return_features = return_features
assert len(n_filters_list) == len(conv1_kernels_size), (
"The length of n_filters_list and conv1_kernel_sizes should be equal."
)
self.ensure_dim = Rearrange("batch chans time -> batch 1 chans time")
self.mstsconv = nn.ModuleList(
[
nn.Sequential(
_TSConv(
self.n_chans,
n_filters_list[b],
conv1_kernels_size[b],
conv2_kernel_size,
depth_multiplier,
pool1_size,
pool2_size,
drop_prob,
activation,
),
Rearrange("batch channels 1 time -> batch time channels"),
)
for b in range(len(n_filters_list))
]
)
branch_linear_in = self._forward_flatten(cat=False)
self.branch_head = nn.ModuleList(
[
_DenseLayers(branch_linear_in[b].shape[1], self.n_outputs)
for b in range(len(n_filters_list))
]
)
seq_len, d_model = self._forward_mstsconv().shape[1:3] # type: ignore
self.transformer = _Transformer(
seq_len,
d_model,
num_heads,
feedforward_ratio,
drop_prob_trans,
num_layers,
)
linear_in = self._forward_flatten().shape[1] # type: ignore
self.flatten_layer = nn.Flatten()
self.final_layer = nn.Linear(linear_in, self.n_outputs)
def _forward_mstsconv(
self, cat: bool = True
) -> Union[torch.Tensor, List[torch.Tensor]]:
x = torch.randn(1, 1, self.n_chans, self.n_times)
x = [tsconv(x) for tsconv in self.mstsconv]
if cat:
x = torch.cat(x, dim=2)
return x
def _forward_flatten(
self, cat: bool = True
) -> Union[torch.Tensor, List[torch.Tensor]]:
x = self._forward_mstsconv(cat)
if cat:
x = self.transformer(x)
x = x.flatten(start_dim=1, end_dim=-1)
else:
x = [_.flatten(start_dim=1, end_dim=-1) for _ in x]
return x
def forward(
self, x: torch.Tensor
) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
# x with shape: (batch, n_chans, n_times)
x = self.ensure_dim(x)
# x with shape: (batch, 1, n_chans, n_times)
x_list = [tsconv(x) for tsconv in self.mstsconv]
# x_list contains 4 tensors, each of shape: [batch_size, seq_len, embed_dim]
branch_preds = [
branch(x_list[idx]) for idx, branch in enumerate(self.branch_head)
]
# branch_preds contains 4 tensors, each of shape: [batch_size, num_classes]
x = torch.stack(x_list, dim=2)
x = x.view(x.size(0), x.size(1), -1)
# x shape after concatenation: [batch_size, seq_len, total_embed_dim]
x = self.transformer(x)
# x shape after transformer: [batch_size, embed_dim]
x = self.final_layer(x)
return (x, branch_preds) if self.return_features else x
Datasets
Models
