spd_learn.models.TSMNet#

class spd_learn.models.TSMNet(n_chans=None, n_temp_filters=4, temp_kernel_length=25, n_spatiotemp_filters=40, n_bimap_filters=20, reeig_threshold=0.0001, n_outputs=None)[source]#

Bases: Module

Tangent Space Mapping Network (TSMNet).

This class implements the TSMNet model [Kobler et al., 2022], which combines a convolutional feature extractor, latent covariance pooling, an SPDNet, and a Tangent Space Mapping (TSM).

TSMNet Architecture

The model consists of the following layers:

  • A convolutional layer that applies the first two layers of ShallowConvNet [Schirrmeister et al., 2017].

  • A sample covariance matrix (SCM) pooling layer that creates SCMs from the convolved signals.

  • An SPDNet that includes a SPDBatchNormMeanVar layer before the LogEig layer.

  • A linear projection head.

Parameters:

  • n_chans (int) – Number of input channels.

  • n_temp_filters (int, default=4) – Number of filters in the temporal convolution layer.

  • temp_kernel_length (int, default=25) – Length of the 1D kernels in the temporal convolution layer.

  • n_spatiotemp_filters (int, default=40) – Number of filters in the spatiotemporal convolution layer.

  • n_bimap_filters (int, default=20) – Number of filters in the BiMap layer.

  • reeig_threshold (float, default=1e-4) – Threshold for the ReEig layer.

  • n_outputs (int) – Number of output dimensions.

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

Forward pass of the TSMNet model.

Parameters:

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

Returns:

Output tensor of shape (batch_size, n_outputs).

Return type:

torch.Tensor