spd_learn.models.TensorCSPNet#

class spd_learn.models.TensorCSPNet(n_chans: int = 22, n_outputs: int = 4, n_patches: int = 4, n_freqs: int = 9, use_mlp: bool = False, tcn_channels: int = 16, dims: Tuple[int, int, int, int] = (22, 36, 36, 22), momentum: float = 0.1)[source]#

Bases: Module

Tensor-CSPNet.

This class implements the Tensor-CSPNet model [Ju and Guan, 2023], an SPDNet framework for EEG-based motor imagery classification.

TensorCSPNet Architecture

The architecture consists of four stages:

  1. Tensor Stacking: Segments the EEG signals into temporospatial-frequency tensors.

  2. Common Spatial Pattern (CSP): Uses modified SPDNet layers to capture spatial patterns.

  3. Temporal Convolution: Captures temporal dynamics using 2D CNN layers.

  4. Classification: Uses a linear layer or a multi-layer perceptron for final classification.

Parameters:

  • n_chans (int, default=22) – Number of input channels.

  • n_outputs (int, default=4) – Number of output classes.

  • n_patches (int, default=4) – Number of patches to split the temporal dimension into.

  • n_freqs (int, default=9) – Number of frequency bands.

  • use_mlp (bool, default=False) – Whether to use a multi-layer perceptron in the final layer.

  • tcn_channels (int, default=16) – Number of channels for the temporal convolutional network.

  • dims (tuple[int, int, int, int], default=(22, 36, 36, 22)) – Dimensions for the BiMap layers.

  • momentum (float, default=0.1) – Momentum factor for the Riemannian Brooks Batch Normalization.

forward(input: Tensor) → Tensor[source]#

Forward pass of the TensorCSPNet model.

Parameters:

input (torch.Tensor) – Input tensor of shape (batch_size, n_freqs, n_chans, n_times).

Returns:

Output tensor of shape (batch_size, n_outputs).

Return type:

torch.Tensor