braindecode 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/models/sparcnet.py

@@ -0,0 +1,378 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from math import floor, log2
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class SPARCNet(EEGModuleMixin, nn.Module):
+    """Seizures, Periodic and Rhythmic pattern Continuum Neural Network (SPaRCNet) from Jing et al. (2023) [jing2023]_.
+
+    This is a temporal CNN model for biosignal classification based on the DenseNet
+    architecture.
+
+    The model is based on the unofficial implementation [Code2023]_.
+
+    .. versionadded:: 0.9
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors.
+
+    Parameters
+    ----------
+    block_layers : int, optional
+        Number of layers per dense block. Default is 4.
+    growth_rate : int, optional
+        Growth rate of the DenseNet. Default is 16.
+    bn_size : int, optional
+        Bottleneck size. Default is 16.
+    drop_prob : float, optional
+        Dropout rate. Default is 0.5.
+    conv_bias : bool, optional
+        Whether to use bias in convolutional layers. Default is True.
+    batch_norm : bool, optional
+        Whether to use batch normalization. Default is True.
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    References
+    ----------
+    .. [jing2023] Jing, J., Ge, W., Hong, S., Fernandes, M. B., Lin, Z.,
+       Yang, C., ... & Westover, M. B. (2023). Development of expert-level
+       classification of seizures and rhythmic and periodic
+       patterns during eeg interpretation. Neurology, 100(17), e1750-e1762.
+    .. [Code2023] Yang, C., Westover, M.B. and Sun, J., 2023. BIOT
+       Biosignal Transformer for Cross-data Learning in the Wild.
+       GitHub https://github.com/ycq091044/BIOT (accessed 2024-02-13)
+
+    """
+
+    def __init__(
+        self,
+        n_chans=None,
+        n_times=None,
+        n_outputs=None,
+        # Neural network parameters
+        block_layers: int = 4,
+        growth_rate: int = 16,
+        bottleneck_size: int = 16,
+        drop_prob: float = 0.5,
+        conv_bias: bool = True,
+        batch_norm: bool = True,
+        activation: nn.Module = nn.ELU,
+        # EEGModuleMixin parameters
+        # (another way to present the same parameters)
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+    ):
+        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
+
+        # add initial convolutional layer
+        # the number of output channels is the smallest power of 2
+        # that is greater than the number of input channels
+        out_channels = 2 ** (floor(log2(self.n_chans)) + 1)
+        first_conv = OrderedDict(
+            [
+                (
+                    "conv0",
+                    nn.Conv1d(
+                        in_channels=self.n_chans,
+                        out_channels=out_channels,
+                        kernel_size=7,
+                        stride=2,
+                        padding=3,
+                        bias=conv_bias,
+                    ),
+                )
+            ]
+        )
+        first_conv["norm0"] = nn.BatchNorm1d(out_channels)
+        first_conv["act_layer"] = activation()
+        first_conv["pool0"] = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+        self.encoder = nn.Sequential(first_conv)
+
+        n_channels = out_channels
+
+        # Adding dense blocks
+        for n_layer in range(floor(log2(self.n_times // 4))):
+            block = _DenseBlock(
+                num_layers=block_layers,
+                in_channels=n_channels,
+                growth_rate=growth_rate,
+                bottleneck_size=bottleneck_size,
+                drop_prob=drop_prob,
+                conv_bias=conv_bias,
+                batch_norm=batch_norm,
+                activation=activation,
+            )
+            self.encoder.add_module("denseblock%d" % (n_layer + 1), block)
+            # update the number of channels after each dense block
+            n_channels = n_channels + block_layers * growth_rate
+
+            trans = _TransitionLayer(
+                in_channels=n_channels,
+                out_channels=n_channels // 2,
+                conv_bias=conv_bias,
+                batch_norm=batch_norm,
+                activation=activation,
+            )
+            self.encoder.add_module("transition%d" % (n_layer + 1), trans)
+            # update the number of channels after each transition layer
+            n_channels = n_channels // 2
+
+        # add final convolutional layer
+        self.final_layer = nn.Sequential(
+            activation(),
+            nn.Linear(n_channels, self.n_outputs),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """
+        Initialize the weights of the model.
+
+        Official init from torch repo, using kaiming_normal for conv layers
+        and normal for linear layers.
+
+        """
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight.data)
+            elif isinstance(m, nn.BatchNorm1d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                m.bias.data.zero_()
+
+    def forward(self, X: torch.Tensor):
+        """
+        Forward pass of the model.
+
+        Parameters
+        ----------
+        X: torch.Tensor
+            The input tensor of the model with shape (batch_size, n_channels, n_times)
+
+        Returns
+        -------
+        torch.Tensor
+            The output tensor of the model with shape (batch_size, n_outputs)
+        """
+        emb = self.encoder(X).squeeze(-1)
+        out = self.final_layer(emb)
+        return out
+
+
+class _DenseLayer(nn.Sequential):
+    """
+    A densely connected layer with batch normalization and dropout.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    growth_rate : int
+        Rate of growth of channels in this layer.
+    bottleneck_size : int
+        Multiplicative factor for the bottleneck layer (does not affect the output size).
+    drop_prob : float, optional
+        Dropout rate. Default is 0.5.
+    conv_bias : bool, optional
+        Whether to use bias in convolutional layers. Default is True.
+    batch_norm : bool, optional
+        Whether to use batch normalization. Default is True.
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    Examples
+    --------
+    >>> x = torch.randn(128, 5, 1000)
+    >>> batch, channels, length = x.shape
+    >>> model = _DenseLayer(channels, 5, 2)
+    >>> y = model(x)
+    >>> y.shape
+    torch.Size([128, 10, 1000])
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        growth_rate: int,
+        bottleneck_size: int,
+        drop_prob: float = 0.5,
+        conv_bias: bool = True,
+        batch_norm: bool = True,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        if batch_norm:
+            self.add_module("norm1", nn.BatchNorm1d(in_channels))
+
+        self.add_module("elu1", activation())
+        self.add_module(
+            "conv1",
+            nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=bottleneck_size * growth_rate,
+                kernel_size=1,
+                stride=1,
+                bias=conv_bias,
+            ),
+        )
+        if batch_norm:
+            self.add_module("norm2", nn.BatchNorm1d(bottleneck_size * growth_rate))
+        self.add_module("elu2", activation())
+        self.add_module(
+            "conv2",
+            nn.Conv1d(
+                in_channels=bottleneck_size * growth_rate,
+                out_channels=growth_rate,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                bias=conv_bias,
+            ),
+        )
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Manually pass through each submodule
+        out = x
+        for layer in self:
+            out = layer(out)
+        # apply dropout using the functional API
+        out = F.dropout(out, p=self.drop_prob, training=self.training)
+        # concatenate input and new features
+        return torch.cat([x, out], dim=1)
+
+
+class _DenseBlock(nn.Sequential):
+    """
+    A densely connected block that uses DenseLayers.
+
+    Parameters
+    ----------
+    num_layers : int
+        Number of layers in this block.
+    in_channels : int
+        Number of input channels.
+    growth_rate : int
+        Rate of growth of channels in this layer.
+    bottleneck_size : int
+        Multiplicative factor for the bottleneck layer (does not affect the output size).
+    drop_prob : float, optional
+        Dropout rate. Default is 0.5.
+    conv_bias : bool, optional
+        Whether to use bias in convolutional layers. Default is True.
+    batch_norm : bool, optional
+        Whether to use batch normalization. Default is True.
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    Examples
+    --------
+    >>> x = torch.randn(128, 5, 1000)
+    >>> batch, channels, length = x.shape
+    >>> model = _DenseBlock(3, channels, 5, 2)
+    >>> y = model(x)
+    >>> y.shape
+    torch.Size([128, 20, 1000])
+    """
+
+    def __init__(
+        self,
+        num_layers,
+        in_channels,
+        growth_rate,
+        bottleneck_size,
+        drop_prob=0.5,
+        conv_bias=True,
+        batch_norm=True,
+        activation: nn.Module = nn.ELU,
+    ):
+        super(_DenseBlock, self).__init__()
+        for idx_layer in range(num_layers):
+            layer = _DenseLayer(
+                in_channels=in_channels + idx_layer * growth_rate,
+                growth_rate=growth_rate,
+                bottleneck_size=bottleneck_size,
+                drop_prob=drop_prob,
+                conv_bias=conv_bias,
+                batch_norm=batch_norm,
+                activation=activation,
+            )
+            self.add_module(f"denselayer{idx_layer + 1}", layer)
+
+
+class _TransitionLayer(nn.Sequential):
+    """
+    A pooling transition layer.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    conv_bias : bool, optional
+        Whether to use bias in convolutional layers. Default is True.
+    batch_norm : bool, optional
+        Whether to use batch normalization. Default is True.
+    activation: nn.Module, default=nn.ELU
+        Activation function class to apply. Should be a PyTorch activation
+        module class like ``nn.ReLU`` or ``nn.ELU``. Default is ``nn.ELU``.
+
+    Examples
+    --------
+    >>> x = torch.randn(128, 5, 1000)
+    >>> model = _TransitionLayer(5, 18)
+    >>> y = model(x)
+    >>> y.shape
+    torch.Size([128, 18, 500])
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        conv_bias=True,
+        batch_norm=True,
+        activation: nn.Module = nn.ELU,
+    ):
+        super(_TransitionLayer, self).__init__()
+        if batch_norm:
+            self.add_module("norm", nn.BatchNorm1d(in_channels))
+        self.add_module("elu", activation())
+        self.add_module(
+            "conv",
+            nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+                stride=1,
+                bias=conv_bias,
+            ),
+        )
+        self.add_module("pool", nn.AvgPool1d(kernel_size=2, stride=2))

braindecode/models/summary.csv

@@ -0,0 +1,41 @@
+Model,Paradigm,Type,Freq(Hz),Hyperparameters,#Parameters,get_#Parameters
+ATCNet,General,Classification,250,"n_chans, n_outputs, n_times",113732,"ATCNet(n_chans=22, n_outputs=4, n_times=1000)"
+AttentionBaseNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",3692,"AttentionBaseNet(n_chans=22, n_outputs=4, n_times=1000)"
+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)"
+BIOT,"Sleep Staging, Epilepsy",Classification,200,"n_chans, n_outputs",3183879,"BIOT(n_chans=2,  n_outputs=5,  n_times=6000)"
+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)"
+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, heads=2, emb_size=16)"
+Deep4Net,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",282879,"Deep4Net(n_chans=22, n_outputs=4, n_times=1000)"
+DeepSleepNet,Sleep Staging,Classification,256,"n_chans, n_outputs",24744837,"DeepSleepNet(n_chans=1, n_outputs=5, n_times=7680, sfreq=256)"
+EEGConformer,General,Classification,250,"n_chans, n_outputs, n_times",789572,"EEGConformer(n_chans=22, n_outputs=4, n_times=1000)."
+EEGInceptionERP,"ERP, SSVEP",Classification,128,"n_chans, n_outputs",14926,"EEGInceptionERP(n_chans=8, n_outputs=2, n_times=128, sfreq=128)"
+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)"
+EEGITNet,Motor Imagery,Classification,125,"n_chans, n_outputs, n_times",5212,"EEGITNet(n_chans=22,  n_outputs=4,  n_times=500)"
+EEGNetv1,General,Classification,128,"n_chans, n_outputs, n_times",3052,"EEGNetv1(n_chans=22,  n_outputs=4,  n_times=512)"
+EEGNetv4,General,Classification,128,"n_chans, n_outputs, n_times",2484,"EEGNetv4(n_chans=22,  n_outputs=4,  n_times=512)"
+EEGNeX,Motor Imagery,Classification,125,"n_chans, n_outputs, n_times",55940,"EEGNeX(n_chans=22,  n_outputs=4,  n_times=500)"
+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)"
+EEGResNet,General,Classification,250,"n_chans, n_outputs, n_times",247484,"EEGResNet(n_chans=22,  n_outputs=4,  n_times=1000)"
+EEGSimpleConv,Motor Imagery,Classification,80,"n_chans, n_outputs, sfreq",730404,"EEGSimpleConv(n_chans=22,  n_outputs=4,  n_times=320, sfreq=80)"
+EEGTCNet,Motor Imagery,Classification,250,"n_chans, n_outputs",4516,"EEGTCNet(n_chans=22,  n_outputs=4,  n_times=1000, kern_length=32)"
+Labram,General,"Classification, Embedding",200,"n_chans, n_outputs, n_times",5866180,"Labram(n_chans=22, n_outputs=4, n_times=1000, sfreq=250)"
+MSVTNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",75494," MSVTNet(n_chans=22,  n_outputs=4,  n_times=1000)"
+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)"
+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])"
+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])"
+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)"
+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)"
+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)"
+ShallowFBCSPNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",46084,"ShallowFBCSPNet(n_chans=22,  n_outputs=4,  n_times=1000, sfreq=250)"
+SleepStagerBlanco2020,Sleep Staging,Classification,100,"n_chans, n_outputs, n_times",2845,"SleepStagerBlanco2020(n_chans=2, n_outputs=5, n_times=3000, sfreq=100)"
+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)"
+SleepStagerEldele2021,Sleep Staging,Classification,100,"n_chans, n_outputs, n_times, sfreq",719925,"SleepStagerEldele2021(n_chans=2, n_outputs=5, n_times=3000, sfreq=100)"
+SPARCNet,Epilepsy,Classification,200,"n_chans, n_outputs, n_times",1141921,"SPARCNet(n_chans=16, n_outputs=6, n_times=2000, sfreq=200)"
+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)"
+TSceptionV1,Emotion Recognition,Classification,256,"n_chans, n_outputs, n_times, sfreq",2187206,"TSceptionV1(n_chans=62, n_outputs=3, n_times=5120, sfreq=256)"
+TIDNet,General,Classification,250,"n_chans, n_outputs, n_times",240404,"TIDNet(n_chans=22,  n_outputs=4,  n_times=1000)"
+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)"
+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)"
+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)"
+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)"
+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)"

braindecode/models/syncnet.py

@@ -0,0 +1,232 @@
+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):
+    """Synchronization Network (SyncNet) from Li, Y et al (2017) [Li2017]_.
+
+    .. 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: 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