spd_learn.models.STaRNet#

class spd_learn.models.STaRNet(n_chans: int, n_outputs: int, spatial_kernel_sizes: Sequence[int] = (3, 11, 22), spatial_expansion_coeff: int = 44, depth_spatial_fusion: int = 8, depth_temporal_map: int = 4, temporal_kernel_sizes: Sequence[int] = (25, 65, 125), bimap_reduced_dimension: int = 48, bimap_dims: Sequence[int] | int | None = None)[source]#

Bases: Module

Spatio-Temporal Riemannian Network (STaRNet).

This class implements the STaRNet [Wang et al., 2024]. STaRNet is a neural network model designed for EEG signal classification, combining multi-scale spatial and temporal convolution branches with a covariance-based SPD feature extraction module.

STaRNet Architecture
Parameters:

  • n_chans (int) – Number of input EEG channels (electrodes).

  • n_outputs (int) – Number of output classes for classification.

  • spatial_kernel_sizes (Sequence[int], optional, default=(3, 11, 22)) – List of kernel sizes for the spatial convolutional layers.

  • spatial_expansion_coeff (int, optional, default=44) – Expansion coefficient for the spatial convolutional layers.

  • depth_spatial_fusion (int, optional, default=8) – Depth of the spatial fusion convolutional layer.

  • depth_temporal_map (int, optional, default=4) – Depth of the temporal mapping convolutional layers.

  • temporal_kernel_sizes (Sequence[int], optional, default=(25, 65, 125)) – List of kernel sizes for the temporal convolutional layers.

  • bimap_reduced_dimension (int, optional, default=48) – Reduced dimension for the BiMap layer in the SPD feature extraction module.

  • bimap_dims (Sequence[int] or int, optional) – Optional alias for bimap_reduced_dimension for API consistency across SPDNet-family models.

Notes

The default dimensions are an engineering configuration for 22-channel EEG. The STaRNet paper does not report every architecture dimension, so these values should not be treated as an exact reproduction of the authors’ model.

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

Apply the network to raw EEG epochs.

Parameters:

x (torch.Tensor) – Input tensor with shape (batch_size, n_chans, n_times).

Returns:

Class logits with shape (batch_size, n_outputs).

Return type:

torch.Tensor