braindecode 1.3.0.dev177069446__py3-none-any.whl

braindecode/models/summary.csv

@@ -0,0 +1,47 @@
+Model,Application,Type,Sampling Frequency (Hz),Hyperparameters,#Parameters,get_#Parameters,Categorization
+ATCNet,General,Classification,250,"n_chans, n_outputs, n_times",113732,"ATCNet(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Recurrent,Attention/Transformer"
+AttentionBaseNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",3692,"AttentionBaseNet(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer"
+BDTCN,Normal Abnormal,Classification,100,"n_chans, n_outputs, n_times",456502,"BDTCN(n_chans=21, n_outputs=2, n_times=6000, n_blocks=5, n_filters=55, kernel_size=16)","Convolution,Recurrent"
+BIOT,"Sleep Staging, Epilepsy",Classification,200,"n_chans, n_outputs",3183879,"BIOT(n_chans=2,  n_outputs=5,  n_times=6000)","Foundation Model"
+ContraWR,Sleep Staging,"Classification, Embedding",125,"n_chans, n_outputs, sfreq",1160165,"ContraWR(n_chans=2, n_outputs=5, n_times=3750, emb_size=256, sfreq=125)",Convolution
+CTNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",26900,"CTNet(n_chans=22, n_outputs=4, n_times=1000, n_filters_time=8, kernel_size=16, num_heads=2, embed_dim=16)","Convolution,Attention/Transformer"
+Deep4Net,General,Classification,250,"n_chans, n_outputs, n_times",282879,"Deep4Net(n_chans=22, n_outputs=4, n_times=1000)","Convolution"
+DeepSleepNet,Sleep Staging,Classification,256,"n_chans, n_outputs",24744837,"DeepSleepNet(n_chans=1, n_outputs=5, n_times=7680, sfreq=256)","Convolution,Recurrent"
+EEGConformer,General,Classification,250,"n_chans, n_outputs, n_times",789572,"EEGConformer(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer"
+EEGInceptionERP,"ERP, SSVEP",Classification,128,"n_chans, n_outputs",14926,"EEGInceptionERP(n_chans=8, n_outputs=2, n_times=128, sfreq=128)","Convolution"
+EEGInceptionMI,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",558028,"EEGInceptionMI(n_chans=22, n_outputs=4, n_times=1000, n_convs=5, n_filters=12)","Convolution"
+EEGITNet,Motor Imagery,Classification,125,"n_chans, n_outputs, n_times",5212,"EEGITNet(n_chans=22,  n_outputs=4,  n_times=500)","Convolution,Recurrent"
+EEGNet,General,Classification,128,"n_chans, n_outputs, n_times",2484,"EEGNet(n_chans=22,  n_outputs=4,  n_times=512)","Convolution"
+EEGNeX,Motor Imagery,Classification,125,"n_chans, n_outputs, n_times",55940,"EEGNeX(n_chans=22,  n_outputs=4,  n_times=500)","Convolution"
+EEGSym,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",299218,"EEGSym(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,Channel"
+EEGMiner,Emotion Recognition,Classification,128,"n_chans, n_outputs, n_times, sfreq",7572,"EEGMiner(n_chans=62, n_outputs=2, n_times=2560, sfreq=128)","Convolution,Interpretability"
+EEGSimpleConv,Motor Imagery,Classification,80,"n_chans, n_outputs, sfreq",730404,"EEGSimpleConv(n_chans=22,  n_outputs=4,  n_times=320, sfreq=80)","Convolution"
+EEGTCNet,Motor Imagery,Classification,250,"n_chans, n_outputs",4516,"EEGTCNet(n_chans=22,  n_outputs=4,  n_times=1000, kern_length=32)","Convolution,Recurrent"
+Labram,General,"Classification, Embedding",200,"n_chans, n_outputs, n_times",5866180,"Labram(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,Foundation Model"
+MSVTNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",75494," MSVTNet(n_chans=22,  n_outputs=4,  n_times=1000)","Convolution,Recurrent,Attention/Transformer"
+SCCNet,Motor Imagery,Classification,125,"n_chans, n_outputs, n_times, sfreq",12070,"SCCNet(n_chans=22,  n_outputs=4,  n_times=1000, sfreq=125)","Convolution"
+SignalJEPA,"Motor Imagery, ERP, SSVEP",Embedding,128,"n_times, chs_info",3456882,"SignalJEPA(n_times=512, chs_info=Lee2019_MI().get_data(subjects=[1])[1]['0']['1train'].info[""chs""][:62])","Convolution,Channel,Foundation Model"
+SignalJEPA_Contextual,"Motor Imagery, ERP, SSVEP",Classification,128,"n_outputs, n_times, chs_info",3459184,"SignalJEPA_Contextual(n_outputs=2, input_window_seconds=4.19, sfreq=128, chs_info=Lee2019_MI().get_data(subjects=[1])[1]['0']['1train'].info[""chs""][:62])","Convolution,Channel,Foundation Model"
+SignalJEPA_PostLocal,"Motor Imagery, ERP, SSVEP",Classification,128,"n_chans, n_outputs, n_times",16142,"SignalJEPA_PostLocal(n_chans=62, n_outputs=2, input_window_seconds=4.19, sfreq=128)","Convolution,Channel,Foundation Model"
+SignalJEPA_PreLocal,"Motor Imagery, ERP, SSVEP",Classification,128,"n_outputs, n_times, chs_info",16142,"SignalJEPA_PreLocal(n_chans=62, n_outputs=2, input_window_seconds=4.19, sfreq=128)","Convolution,Channel,Foundation Model"
+SincShallowNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",21892,"SincShallowNet(n_chans=22,  n_outputs=4,  n_times=1000, sfreq=250)","Convolution,Interpretability"
+ShallowFBCSPNet,General,Classification,250,"n_chans, n_outputs, n_times",46084,"ShallowFBCSPNet(n_chans=22,  n_outputs=4,  n_times=1000, sfreq=250)","Convolution"
+SleepStagerBlanco2020,Sleep Staging,Classification,100,"n_chans, n_outputs, n_times",2845,"SleepStagerBlanco2020(n_chans=2, n_outputs=5, n_times=3000, sfreq=100)","Convolution"
+SleepStagerChambon2018,Sleep Staging,Classification,128,"n_chans, n_outputs, n_times, sfreq",5835,"SleepStagerChambon2018(n_chans=2, n_outputs=5, n_times=3840, sfreq=128)","Convolution"
+AttnSleep,Sleep Staging,Classification,100,"n_chans, n_outputs, n_times, sfreq",719925,"AttnSleep(n_chans=2, n_outputs=5, n_times=3000, sfreq=100)","Convolution, Attention/Transformer"
+SPARCNet,Epilepsy,Classification,200,"n_chans, n_outputs, n_times",1141921,"SPARCNet(n_chans=16, n_outputs=6, n_times=2000, sfreq=200)","Convolution"
+SyncNet,"Emotion Recognition, Alcoholism",Classification,256,"n_chans, n_outputs, n_times",554,"SyncNet(n_chans=62, n_outputs=3, n_times=5120, sfreq=256)","Interpretability"
+TSception,Emotion Recognition,Classification,256,"n_chans, n_outputs, n_times, sfreq",2187206,"TSception(n_chans=62, n_outputs=3, n_times=5120, sfreq=256)","Convolution"
+TIDNet,General,Classification,250,"n_chans, n_outputs, n_times",240404,"TIDNet(n_chans=22,  n_outputs=4,  n_times=1000)","Convolution"
+USleep,Sleep Staging,Classification,128,"n_chans, n_outputs, n_times, sfreq",2482011,"USleep(n_chans=2, n_outputs=5, n_times=3000, sfreq=100)","Convolution"
+FBCNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",11812,"FCNet(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,FilterBank"
+FBMSNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",16231,"FBMSNet(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,FilterBank"
+FBLightConvNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",6596,"FBLightConvNet(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,FilterBank"
+IFNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times, sfreq",9860,"IFNet(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Convolution,FilterBank"
+BrainModule,Speech Decoding,Classification,250,"n_chans, n_outputs, n_times, sfreq",6186909,"BrainModule(n_chans=64, n_outputs=29, n_times=160, sfreq=1000)","Convolution"
+PBT,General,Classification,250,"n_chans, n_outputs, n_times",818948,"PBT(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)","Foundation Model"
+SSTDPN,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",19502,"SSTDPN(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer"
+BENDR,General,"Classification,Embedding",250,"n_chans, n_times, n_outputs",157141049,"BENDR(n_chans=22, n_outputs=4, n_times=1000)","Foundation Model,Convolution"
+LUNA,General,"Classification,Embedding",128,"n_chans, n_times, sfreq, chs_info",7100731,"LUNA(n_chans=22, n_times=512, sfreq=128)","Convolution,Channel,Foundation Model"
+MEDFormer,General,Classification,250,"n_chans, n_outputs, n_times",5313924,"MEDFormer(n_chans=22, n_outputs=4, n_times=1000)","Foundation Model,Convolution"
+REVE,General,Classification,200,"n_outputs, n_times, n_chans",69481476,"REVE(n_times=1000, n_outputs=4, n_chans=19)","Foundation Model,Attention/Transformer"

braindecode/models/syncnet.py

@@ -0,0 +1,234 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.layers.torch import Rearrange
+from numpy import arange, ceil
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SyncNet(EEGModuleMixin, nn.Module):
+    r"""Synchronization Network (SyncNet) from Li, Y et al (2017) [Li2017]_.
+
+    :bdg-warning:`Interpretability`
+
+    .. figure:: https://braindecode.org/dev/_static/model/SyncNet.png
+        :align: center
+        :alt: SyncNet Architecture
+
+    SyncNet uses parameterized 1-dimensional convolutional filters inspired by
+    the Morlet wavelet to extract features from EEG signals. The filters are
+    dynamically generated based on learnable parameters that control the
+    oscillation and decay characteristics.
+
+    The filter for channel ``c`` and filter ``k`` is defined as:
+
+    .. math::
+
+        f_c^{(k)}(\\tau) = amplitude_c^{(k)} \\cos(\\omega^{(k)} \\tau + \\phi_c^{(k)}) \\exp(-\\beta^{(k)} \\tau^2)
+
+    where:
+    - :math:`amplitude_c^{(k)}` is the amplitude parameter (channel-specific).
+    - :math:`\\omega^{(k)}` is the frequency parameter (shared across channels).
+    - :math:`\\phi_c^{(k)}` is the phase shift (channel-specific).
+    - :math:`\\beta^{(k)}` is the decay parameter (shared across channels).
+    - :math:`\\tau` is the time index.
+
+    Parameters
+    ----------
+    num_filters : int, optional
+        Number of filters in the convolutional layer. Default is 1.
+    filter_width : int, optional
+        Width of the convolutional filters. Default is 40.
+    pool_size : int, optional
+        Size of the pooling window. Default is 40.
+    activation : nn.Module, optional
+        Activation function to apply after pooling. Default is ``nn.ReLU``.
+    ampli_init_values : tuple of float, optional
+        The initialization range for amplitude parameter using uniform
+        distribution. Default is (-0.05, 0.05).
+    omega_init_values : tuple of float, optional
+        The initialization range for omega parameters using uniform
+        distribution. Default is (0, 1).
+    beta_init_values : tuple of float, optional
+        The initialization range for beta parameters using uniform
+        distribution. Default is (0, 1). Default is (0, 0.05).
+    phase_init_values : tuple of float, optional
+        The initialization range for phase parameters using `normal`
+        distribution. Default is (0, 1). Default is (0, 0.05).
+
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct! it has not been checked
+    by original authors. The modifications are based on derivated code from
+    [CodeICASSP2025]_.
+
+
+    References
+    ----------
+    .. [Li2017] Li, Y., Dzirasa, K., Carin, L., & Carlson, D. E. (2017).
+       Targeting EEG/LFP synchrony with neural nets. Advances in neural
+       information processing systems, 30.
+    .. [CodeICASSP2025] Code from Baselines for EEG-Music Emotion Recognition
+       Grand Challenge at ICASSP 2025.
+       https://github.com/SalvoCalcagno/eeg-music-challenge-icassp-2025-baselines
+
+    """
+
+    def __init__(
+        self,
+        # braindecode convention
+        n_chans=None,
+        n_times=None,
+        n_outputs=None,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # model parameters
+        num_filters=1,
+        filter_width=40,
+        pool_size=40,
+        activation: type[nn.Module] = nn.ReLU,
+        ampli_init_values: tuple[float, float] = (-0.05, 0.05),
+        omega_init_values: tuple[float, float] = (0.0, 1.0),
+        beta_init_values: tuple[float, float] = (0.0, 0.05),
+        phase_init_values: tuple[float, float] = (0.0, 0.05),
+    ):
+        super().__init__(
+            n_chans=n_chans,
+            n_times=n_times,
+            n_outputs=n_outputs,
+            chs_info=chs_info,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+
+        self.num_filters = num_filters
+        self.filter_width = filter_width
+        self.pool_size = pool_size
+        self.activation = activation()
+        self.ampli_init_values = ampli_init_values
+        self.omega_init_values = omega_init_values
+        self.beta_init_values = beta_init_values
+        self.phase_init_values = phase_init_values
+
+        # Initialize parameters
+        self.amplitude = nn.Parameter(
+            torch.FloatTensor(1, 1, self.n_chans, self.num_filters).uniform_(
+                self.ampli_init_values[0], self.ampli_init_values[1]
+            )
+        )
+        self.omega = nn.Parameter(
+            torch.FloatTensor(1, 1, 1, self.num_filters).uniform_(
+                self.omega_init_values[0], self.omega_init_values[1]
+            )
+        )
+
+        self.bias = nn.Parameter(torch.zeros(self.num_filters))
+
+        # Calculate the output size after pooling
+        self.classifier_input_size = int(
+            ceil(float(self.n_times) / float(self.pool_size)) * self.num_filters
+        )
+
+        # Create time vector t
+        if self.filter_width % 2 == 0:
+            t_range = arange(-int(self.filter_width / 2), int(self.filter_width / 2))
+        else:
+            t_range = arange(
+                -int((self.filter_width - 1) / 2), int((self.filter_width - 1) / 2) + 1
+            )
+
+        t_np = t_range.reshape(1, self.filter_width, 1, 1)
+        self.t = nn.Parameter(torch.FloatTensor(t_np))
+        # Phase Shift
+        self.phi_ini = nn.Parameter(
+            torch.FloatTensor(1, 1, self.n_chans, self.num_filters).normal_(
+                self.beta_init_values[0], self.beta_init_values[1]
+            )
+        )
+        self.beta = nn.Parameter(
+            torch.FloatTensor(1, 1, 1, self.num_filters).uniform_(
+                self.phase_init_values[0], self.phase_init_values[1]
+            )
+        )
+
+        self.padding = self._compute_padding(filter_width=self.filter_width)
+        self.pad_input = nn.ConstantPad1d(self.padding, 0.0)
+        self.pad_res = nn.ConstantPad1d(self.padding, 0.0)
+
+        # Define pooling and classifier layers
+        self.pool = nn.MaxPool2d((1, self.pool_size), stride=(1, self.pool_size))
+
+        self.ensuredim = Rearrange("batch ch time -> batch ch 1 time")
+
+        self.final_layer = nn.Linear(self.classifier_input_size, self.n_outputs)
+
+    def forward(self, x):
+        """Forward pass of the SyncNet model.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (batch_size, n_chans, n_times)
+
+        Returns
+        -------
+        out : torch.Tensor
+            Output tensor of shape (batch_size, n_outputs).
+
+        """
+        # Ensure input tensor has shape (batch_size, n_chans, 1, n_times)
+        x = self.ensuredim(x)
+        # Output: (batch_size, n_chans, 1, n_times)
+
+        # Compute the oscillatory component
+        W_osc = self.amplitude * torch.cos(self.t * self.omega + self.phi_ini)
+        # W_osc is (1, filter_width, n_chans, 1)
+
+        # Compute the decay component
+        t_squared = torch.pow(self.t, 2)  # Shape: (filter_width,)
+        t_squared_beta = t_squared * self.beta  # Shape: (filter_width, num_filters)
+        W_decay = torch.exp(-t_squared_beta)
+        # W_osc is (1, filter_width, 1, 1)
+
+        # Combine oscillatory and decay components
+        # W shape: (1, n_chans, num_filters, filter_width)
+        W = W_osc * W_decay
+        # W shape will be: (1, filter_width, n_chans, 1)
+
+        W = W.view(self.num_filters, self.n_chans, 1, self.filter_width)
+
+        # Apply convolution
+        x_padded = self.pad_input(x.float())
+
+        res = F.conv2d(x_padded, W.float(), bias=self.bias, stride=1)
+
+        # Apply padding to the convolution result
+        res_padded = self.pad_res(res)
+        res_pooled = self.pool(res_padded)
+
+        # Flatten the result
+        res_flat = res_pooled.view(-1, self.classifier_input_size)
+
+        # Ensure beta remains non-negative
+        self.beta.data.clamp_(min=0)
+
+        # Apply activation
+        out = self.activation(res_flat)
+        # Apply classifier
+        out = self.final_layer(out)
+
+        return out
+
+    @staticmethod
+    def _compute_padding(filter_width):
+        # Compute padding
+        P = filter_width - 2
+        if P % 2 == 0:
+            padding = (P // 2, P // 2 + 1)
+        else:
+            padding = (P // 2, P // 2)
+        return padding

braindecode/models/tcn.py

@@ -0,0 +1,275 @@
+# Authors: Patryk Chrabaszcz
+#          Lukas Gemein <l.gemein@gmail.com>
+#
+# License: BSD-3
+import torch
+from torch import nn
+from torch.nn import init
+from torch.nn.utils.parametrizations import weight_norm
+
+from braindecode.models.base import EEGModuleMixin
+from braindecode.modules import Chomp1d, Ensure4d, SqueezeFinalOutput
+
+
+class BDTCN(EEGModuleMixin, nn.Module):
+    r"""Braindecode TCN from Gemein, L et al (2020) [gemein2020]_.
+
+    :bdg-success:`Convolution` :bdg-secondary:`Recurrent`
+
+    .. figure:: https://ars.els-cdn.com/content/image/1-s2.0-S1053811920305073-gr3_lrg.jpg
+       :align: center
+       :alt: Braindecode TCN Architecture
+
+    See [gemein2020]_ for details.
+
+    Parameters
+    ----------
+    n_filters: int
+        number of output filters of each convolution
+    n_blocks: int
+        number of temporal blocks in the network
+    kernel_size: int
+        kernel size of the convolutions
+    drop_prob: float
+        dropout probability
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+
+    References
+    ----------
+    .. [gemein2020] Gemein, L. A., Schirrmeister, R. T., Chrabąszcz, P., Wilson, D.,
+       Boedecker, J., Schulze-Bonhage, A., ... & Ball, T. (2020). Machine-learning-based
+       diagnostics of EEG pathology. NeuroImage, 220, 117021.
+    """
+
+    def __init__(
+        self,
+        # Braindecode parameters
+        n_chans=None,
+        n_outputs=None,
+        chs_info=None,
+        n_times=None,
+        sfreq=None,
+        input_window_seconds=None,
+        # Model's parameters
+        n_blocks=3,
+        n_filters=30,
+        kernel_size=5,
+        drop_prob=0.5,
+        activation: type[nn.Module] = nn.ReLU,
+    ):
+        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, sfreq, input_window_seconds
+
+        self.base_tcn = TCN(
+            n_chans=self.n_chans,
+            n_outputs=self.n_outputs,
+            n_blocks=n_blocks,
+            n_filters=n_filters,
+            kernel_size=kernel_size,
+            drop_prob=drop_prob,
+            activation=activation,
+        )
+
+        self.final_layer = torch.nn.Sequential(
+            torch.nn.AdaptiveAvgPool1d(1), torch.nn.Flatten()
+        )
+
+    def forward(self, x):
+        x = self.base_tcn(x)
+        x = self.final_layer(x)
+        return x
+
+
+class TCN(nn.Module):
+    r"""Temporal Convolutional Network (TCN) from Bai et al. 2018 [Bai2018]_.
+
+    See [Bai2018]_ for details.
+
+    Code adapted from https://github.com/locuslab/TCN/blob/master/TCN/tcn.py
+
+    Parameters
+    ----------
+    n_filters: int
+        number of output filters of each convolution
+    n_blocks: int
+        number of temporal blocks in the network
+    kernel_size: int
+        kernel size of the convolutions
+    drop_prob: float
+        dropout probability
+    activation: nn.Module, default=nn.ReLU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ReLU``.
+
+    References
+    ----------
+    .. [Bai2018] Bai, S., Kolter, J. Z., & Koltun, V. (2018).
+       An empirical evaluation of generic convolutional and recurrent networks
+       for sequence modeling.
+       arXiv preprint arXiv:1803.01271.
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_outputs=None,
+        n_blocks=3,
+        n_filters=30,
+        kernel_size=5,
+        drop_prob=0.5,
+        activation: type[nn.Module] = nn.ReLU,
+    ):
+        super().__init__()
+        self.mapping = {
+            "fc.weight": "final_layer.fc.weight",
+            "fc.bias": "final_layer.fc.bias",
+        }
+        self.ensuredims = Ensure4d()
+        t_blocks = nn.Sequential()
+        for i in range(n_blocks):
+            n_inputs = n_chans if i == 0 else n_filters
+            dilation_size = 2**i
+            t_blocks.add_module(
+                "temporal_block_{:d}".format(i),
+                _TemporalBlock(
+                    n_inputs=n_inputs,
+                    n_outputs=n_filters,
+                    kernel_size=kernel_size,
+                    stride=1,
+                    dilation=dilation_size,
+                    padding=(kernel_size - 1) * dilation_size,
+                    drop_prob=drop_prob,
+                    activation=activation,
+                ),
+            )
+        self.temporal_blocks = t_blocks
+
+        self.final_layer = _FinalLayer(
+            in_features=n_filters,
+            out_features=n_outputs,
+        )
+        self.min_len = 1
+        for i in range(n_blocks):
+            dilation = 2**i
+            self.min_len += 2 * (kernel_size - 1) * dilation
+
+        # start in eval mode
+        self.train()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Batch of EEG windows of shape (batch_size, n_channels, n_times).
+        """
+        x = self.ensuredims(x)
+        # x is in format: B x C x T x 1
+        (batch_size, _, time_size, _) = x.size()
+        assert time_size >= self.min_len
+        # remove empty trailing dimension
+        x = x.squeeze(3)
+        x = self.temporal_blocks(x)
+        # Convert to: B x T x C
+        x = x.transpose(1, 2).contiguous()
+
+        out = self.final_layer(x, batch_size, time_size, self.min_len)
+
+        return out
+
+
+class _FinalLayer(nn.Module):
+    def __init__(self, in_features, out_features):
+        super().__init__()
+
+        self.fc = nn.Linear(in_features=in_features, out_features=out_features)
+
+        self.out_fun = nn.Identity()
+
+        self.squeeze = SqueezeFinalOutput()
+
+    def forward(
+        self, x: torch.Tensor, batch_size: int, time_size: int, min_len: int
+    ) -> torch.Tensor:
+        fc_out = self.fc(x.view(batch_size * time_size, x.size(2)))
+        fc_out = self.out_fun(fc_out)
+        fc_out = fc_out.view(batch_size, time_size, fc_out.size(1))
+
+        out_size = 1 + max(0, time_size - min_len)
+        out = fc_out[:, -out_size:, :].transpose(1, 2)
+        # re-add 4th dimension for compatibility with braindecode
+        return self.squeeze(out[:, :, :, None])
+
+
+class _TemporalBlock(nn.Module):
+    def __init__(
+        self,
+        n_inputs,
+        n_outputs,
+        kernel_size,
+        stride,
+        dilation,
+        padding,
+        drop_prob,
+        activation: type[nn.Module] = nn.ReLU,
+    ):
+        super().__init__()
+        self.conv1 = weight_norm(
+            nn.Conv1d(
+                n_inputs,
+                n_outputs,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+            )
+        )
+        self.chomp1 = Chomp1d(padding)
+        self.relu1 = activation()
+        self.dropout1 = nn.Dropout2d(drop_prob)
+
+        self.conv2 = weight_norm(
+            nn.Conv1d(
+                n_outputs,
+                n_outputs,
+                kernel_size,
+                stride=stride,
+                padding=padding,
+                dilation=dilation,
+            )
+        )
+        self.chomp2 = Chomp1d(padding)
+        self.relu2 = activation()
+        self.dropout2 = nn.Dropout2d(drop_prob)
+
+        self.downsample = (
+            nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
+        )
+        self.relu = activation()
+
+        init.normal_(self.conv1.weight, 0, 0.01)
+        init.normal_(self.conv2.weight, 0, 0.01)
+        if self.downsample is not None:
+            init.normal_(self.downsample.weight, 0, 0.01)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.chomp1(out)
+        out = self.relu1(out)
+        out = self.dropout1(out)
+        out = self.conv2(out)
+        out = self.chomp2(out)
+        out = self.relu2(out)
+        out = self.dropout2(out)
+        res = x if self.downsample is None else self.downsample(x)
+        return self.relu(out + res)