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

braindecode/models/usleep.py

@@ -3,122 +3,33 @@
 #
 # License: BSD (3-clause)
 
+
 import numpy as np
 import torch
 from torch import nn
-from .base import EEGModuleMixin, deprecated_args
-
-
-def _crop_tensors_to_match(x1, x2, axis=-1):
-    """Crops two tensors to their lowest-common-dimension along an axis."""
-    dim_cropped = min(x1.shape[axis], x2.shape[axis])
-
-    x1_cropped = torch.index_select(
-        x1, dim=axis,
-        index=torch.arange(dim_cropped).to(device=x1.device)
-    )
-    x2_cropped = torch.index_select(
-        x2, dim=axis,
-        index=torch.arange(dim_cropped).to(device=x1.device)
-    )
-    return x1_cropped, x2_cropped
-
-
-class _EncoderBlock(nn.Module):
-    """Encoding block for a timeseries x of shape (B, C, T)."""
-
-    def __init__(self,
-                 in_channels=2,
-                 out_channels=2,
-                 kernel_size=9,
-                 downsample=2):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.kernel_size = kernel_size
-        self.downsample = downsample
-
-        self.block_prepool = nn.Sequential(
-            nn.Conv1d(in_channels=in_channels,
-                      out_channels=out_channels,
-                      kernel_size=kernel_size,
-                      padding='same'),
-            nn.ELU(),
-            nn.BatchNorm1d(num_features=out_channels),
-        )
-
-        self.pad = nn.ConstantPad1d(padding=1, value=0)
-        self.maxpool = nn.MaxPool1d(
-            kernel_size=self.downsample, stride=self.downsample)
-
-    def forward(self, x):
-        x = self.block_prepool(x)
-        residual = x
-        if x.shape[-1] % 2:
-            x = self.pad(x)
-        x = self.maxpool(x)
-        return x, residual
-
-
-class _DecoderBlock(nn.Module):
-    """Decoding block for a timeseries x of shape (B, C, T)."""
-
-    def __init__(self,
-                 in_channels=2,
-                 out_channels=2,
-                 kernel_size=9,
-                 upsample=2,
-                 with_skip_connection=True):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.kernel_size = kernel_size
-        self.upsample = upsample
-        self.with_skip_connection = with_skip_connection
-
-        self.block_preskip = nn.Sequential(
-            nn.Upsample(scale_factor=upsample),
-            nn.Conv1d(in_channels=in_channels,
-                      out_channels=out_channels,
-                      kernel_size=2,
-                      padding='same'),
-            nn.ELU(),
-            nn.BatchNorm1d(num_features=out_channels),
-        )
-        self.block_postskip = nn.Sequential(
-            nn.Conv1d(
-                in_channels=(
-                    2 * out_channels if with_skip_connection else out_channels),
-                out_channels=out_channels,
-                kernel_size=kernel_size,
-                padding='same'),
-            nn.ELU(),
-            nn.BatchNorm1d(num_features=out_channels),
-        )
 
-    def forward(self, x, residual):
-        x = self.block_preskip(x)
-        if self.with_skip_connection:
-            x, residual = _crop_tensors_to_match(x, residual, axis=-1)  # in case of mismatch
-            x = torch.cat([x, residual], axis=1)  # (B, 2 * C, T)
-        x = self.block_postskip(x)
-        return x
+from braindecode.models.base import EEGModuleMixin
 
 
 class USleep(EEGModuleMixin, nn.Module):
-    """Sleep staging architecture from Perslev et al 2021.
+    """
+    Sleep staging architecture from Perslev et al. (2021) [1]_.
+
+    .. figure:: https://media.springernature.com/full/springer-static/image/art%3A10.1038%2Fs41746-021-00440-5/MediaObjects/41746_2021_440_Fig2_HTML.png
+        :align: center
+        :alt: USleep Architecture
 
     U-Net (autoencoder with skip connections) feature-extractor for sleep
     staging described in [1]_.
 
     For the encoder ('down'):
-        -- the temporal dimension shrinks (via maxpooling in the time-domain)
-        -- the spatial dimension expands (via more conv1d filters in the
-           time-domain)
+        - the temporal dimension shrinks (via maxpooling in the time-domain)
+        - the spatial dimension expands (via more conv1d filters in the time-domain)
+
     For the decoder ('up'):
-        -- the temporal dimension expands (via upsampling in the time-domain)
-        -- the spatial dimension shrinks (via fewer conv1d filters in the
-           time-domain)
+        - the temporal dimension expands (via upsampling in the time-domain)
+        - the spatial dimension shrinks (via fewer conv1d filters in the time-domain)
+
     Both do so at exponential rates.
 
     Parameters
@@ -128,12 +39,12 @@ class USleep(EEGModuleMixin, nn.Module):
     sfreq : float
         EEG sampling frequency. Set to 128 in [1]_.
     depth : int
-        Number of conv blocks in encoding layer (number of 2x2 max pools)
-        Note: each block halve the spatial dimensions of the features.
+        Number of conv blocks in encoding layer (number of 2x2 max pools).
+        Note: each block halves the spatial dimensions of the features.
     n_time_filters : int
         Initial number of convolutional filters. Set to 5 in [1]_.
     complexity_factor : float
-        Multiplicative factor for number of channels at each layer of the U-Net.
+        Multiplicative factor for the number of channels at each layer of the U-Net.
         Set to 2 in [1]_.
     with_skip_connection : bool
         If True, use skip connections in decoder blocks.
@@ -147,47 +58,35 @@ class USleep(EEGModuleMixin, nn.Module):
     ensure_odd_conv_size : bool
         If True and the size of the convolutional kernel is an even number, one
         will be added to it to ensure it is odd, so that the decoder blocks can
-        work. This can ne useful when using different sampling rates from 128
+        work. This can be useful when using different sampling rates from 128
         or 100 Hz.
-    in_chans : int
-        Alias for n_chans.
-    n_classes : int
-        Alias for n_outputs.
-    input_size_s : float
-        Alias for input_window_seconds.
+    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
     ----------
     .. [1] Perslev M, Darkner S, Kempfner L, Nikolic M, Jennum PJ, Igel C.
-           U-Sleep: resilient high-frequency sleep staging. npj Digit. Med. 4, 72 (2021).
-           https://github.com/perslev/U-Time/blob/master/utime/models/usleep.py
+       U-Sleep: resilient high-frequency sleep staging. *npj Digit. Med.* 4, 72 (2021).
+       https://github.com/perslev/U-Time/blob/master/utime/models/usleep.py
     """
 
     def __init__(
-            self,
-            n_chans=2,
-            sfreq=128,
-            depth=12,
-            n_time_filters=5,
-            complexity_factor=1.67,
-            with_skip_connection=True,
-            n_outputs=5,
-            input_window_seconds=30,
-            time_conv_size_s=9 / 128,
-            ensure_odd_conv_size=False,
-            chs_info=None,
-            n_times=None,
-            in_chans=None,
-            n_classes=None,
-            input_size_s=None,
-            add_log_softmax=False,
+        self,
+        n_chans=None,
+        sfreq=None,
+        depth=12,
+        n_time_filters=5,
+        complexity_factor=1.67,
+        with_skip_connection=True,
+        n_outputs=5,
+        input_window_seconds=None,
+        time_conv_size_s=9 / 128,
+        ensure_odd_conv_size=False,
+        activation: nn.Module = nn.ELU,
+        chs_info=None,
+        n_times=None,
     ):
-        n_chans, n_outputs, input_window_seconds = deprecated_args(
-            self,
-            ("in_chans", "n_chans", in_chans, n_chans),
-            ("n_classes", "n_outputs", n_classes, n_outputs),
-            ("input_size_s", "input_window_seconds", input_size_s, input_window_seconds),
-        )
         super().__init__(
             n_outputs=n_outputs,
             n_chans=n_chans,
@@ -195,16 +94,14 @@ class USleep(EEGModuleMixin, nn.Module):
             n_times=n_times,
             input_window_seconds=input_window_seconds,
             sfreq=sfreq,
-            add_log_softmax=add_log_softmax,
         )
         del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
-        del in_chans, n_classes, input_size_s
 
         self.mapping = {
-            'clf.3.weight': 'final_layer.0.weight',
-            'clf.3.bias': 'final_layer.0.bias',
-            'clf.5.weight': 'final_layer.2.weight',
-            'clf.5.bias': 'final_layer.2.bias'
+            "clf.3.weight": "final_layer.0.weight",
+            "clf.3.bias": "final_layer.0.bias",
+            "clf.5.weight": "final_layer.2.weight",
+            "clf.5.bias": "final_layer.2.bias",
         }
 
         max_pool_size = 2  # Hardcoded to avoid dimensional errors
@@ -214,8 +111,9 @@ class USleep(EEGModuleMixin, nn.Module):
                 time_conv_size += 1
             else:
                 raise ValueError(
-                    'time_conv_size must be an odd number to accommodate the '
-                    'upsampling step in the decoder blocks.')
+                    "time_conv_size must be an odd number to accommodate the "
+                    "upsampling step in the decoder blocks."
+                )
 
         channels = [self.n_chans]
         n_filters = n_time_filters
@@ -225,38 +123,42 @@ class USleep(EEGModuleMixin, nn.Module):
         self.channels = channels
 
         # Instantiate encoder
-        encoder = list()
-        for idx in range(depth):
-            encoder += [
-                _EncoderBlock(in_channels=channels[idx],
-                              out_channels=channels[idx + 1],
-                              kernel_size=time_conv_size,
-                              downsample=max_pool_size)
-            ]
-        self.encoder = nn.Sequential(*encoder)
+        self.encoder_blocks = nn.ModuleList(
+            _EncoderBlock(
+                in_channels=channels[idx],
+                out_channels=channels[idx + 1],
+                kernel_size=time_conv_size,
+                downsample=max_pool_size,
+                activation=activation,
+            )
+            for idx in range(depth)
+        )
 
         # Instantiate bottom (channels increase, temporal dim stays the same)
         self.bottom = nn.Sequential(
-            nn.Conv1d(in_channels=channels[-2],
-                      out_channels=channels[-1],
-                      kernel_size=time_conv_size,
-                      padding=(time_conv_size - 1) // 2),  # preserves dimension
-            nn.ELU(),
+            nn.Conv1d(
+                in_channels=channels[-2],
+                out_channels=channels[-1],
+                kernel_size=time_conv_size,
+                padding=(time_conv_size - 1) // 2,
+            ),  # preserves dimension
+            activation(),
             nn.BatchNorm1d(num_features=channels[-1]),
         )
 
         # Instantiate decoder
-        decoder = list()
         channels_reverse = channels[::-1]
-        for idx in range(depth):
-            decoder += [
-                _DecoderBlock(in_channels=channels_reverse[idx],
-                              out_channels=channels_reverse[idx + 1],
-                              kernel_size=time_conv_size,
-                              upsample=max_pool_size,
-                              with_skip_connection=with_skip_connection)
-            ]
-        self.decoder = nn.Sequential(*decoder)
+        self.decoder_blocks = nn.ModuleList(
+            _DecoderBlock(
+                in_channels=channels_reverse[idx],
+                out_channels=channels_reverse[idx + 1],
+                kernel_size=time_conv_size,
+                upsample=max_pool_size,
+                with_skip_connection=with_skip_connection,
+                activation=activation,
+            )
+            for idx in range(depth)
+        )
 
         # The temporal dimension remains unchanged
         # (except through the AvgPooling which collapses it to 1)
@@ -282,7 +184,7 @@ class USleep(EEGModuleMixin, nn.Module):
                 stride=1,
                 padding=0,
             ),  # output is (B, n_classes, S)
-            nn.ELU(),
+            activation(),
             nn.Conv1d(
                 in_channels=self.n_outputs,
                 out_channels=self.n_outputs,
@@ -290,11 +192,11 @@ class USleep(EEGModuleMixin, nn.Module):
                 stride=1,
                 padding=0,
             ),
-            nn.LogSoftmax(dim=1) if self.add_log_softmax else nn.Identity(),
+            nn.Identity(),
             # output is (B, n_classes, S)
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """If input x has shape (B, S, C, T), return y_pred of shape (B, n_classes, S).
         If input x has shape (B, C, T), return y_pred of shape (B, n_classes).
         """
@@ -305,7 +207,7 @@ class USleep(EEGModuleMixin, nn.Module):
 
         # encoder
         residuals = []
-        for down in self.encoder:
+        for down in self.encoder_blocks:
             x, res = down(x)
             residuals.append(res)
 
@@ -313,9 +215,11 @@ class USleep(EEGModuleMixin, nn.Module):
         x = self.bottom(x)
 
         # decoder
-        residuals = residuals[::-1]  # flip order
-        for up, res in zip(self.decoder, residuals):
-            x = up(x, res)
+        num_blocks = len(self.decoder_blocks)  # statically known
+        for idx, dec in enumerate(self.decoder_blocks):
+            # pick the matching residual in reverse order
+            res = residuals[num_blocks - 1 - idx]
+            x = dec(x, res)
 
         # classifier
         x = self.clf(x)
@@ -325,3 +229,112 @@ class USleep(EEGModuleMixin, nn.Module):
             y_pred = y_pred[:, :, 0]
 
         return y_pred
+
+
+class _EncoderBlock(nn.Module):
+    """Encoding block for a timeseries x of shape (B, C, T)."""
+
+    def __init__(
+        self,
+        in_channels=2,
+        out_channels=2,
+        kernel_size=9,
+        downsample=2,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.downsample = downsample
+
+        self.block_prepool = nn.Sequential(
+            nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                padding="same",
+            ),
+            activation(),
+            nn.BatchNorm1d(num_features=out_channels),
+        )
+
+        self.pad = nn.ConstantPad1d(padding=1, value=0.0)
+        self.maxpool = nn.MaxPool1d(kernel_size=self.downsample, stride=self.downsample)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = self.block_prepool(x)
+        residual = x
+        if x.shape[-1] % 2:
+            x = self.pad(x)
+        x = self.maxpool(x)
+        return x, residual
+
+
+class _DecoderBlock(nn.Module):
+    """Decoding block for a timeseries x of shape (B, C, T)."""
+
+    def __init__(
+        self,
+        in_channels=2,
+        out_channels=2,
+        kernel_size=9,
+        upsample=2,
+        with_skip_connection=True,
+        activation: nn.Module = nn.ELU,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.upsample = upsample
+        self.with_skip_connection = with_skip_connection
+
+        self.block_preskip = nn.Sequential(
+            nn.Upsample(scale_factor=upsample),
+            nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=2,
+                padding="same",
+            ),
+            activation(),
+            nn.BatchNorm1d(num_features=out_channels),
+        )
+        self.block_postskip = nn.Sequential(
+            nn.Conv1d(
+                in_channels=(
+                    2 * out_channels if with_skip_connection else out_channels
+                ),
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                padding="same",
+            ),
+            activation(),
+            nn.BatchNorm1d(num_features=out_channels),
+        )
+
+    def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
+        x = self.block_preskip(x)
+        if self.with_skip_connection:
+            x, residual = self._crop_tensors_to_match(
+                x, residual, axis=-1
+            )  # in case of mismatch
+            x = torch.cat([x, residual], dim=1)  # (B, 2 * C, T)
+        x = self.block_postskip(x)
+        return x
+
+    @staticmethod
+    def _crop_tensors_to_match(
+        x1: torch.Tensor, x2: torch.Tensor, axis: int = -1
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Crops two tensors to their lowest-common-dimension along an axis."""
+        dim_cropped = min(x1.shape[axis], x2.shape[axis])
+
+        x1_cropped = torch.index_select(
+            x1, dim=axis, index=torch.arange(dim_cropped).to(device=x1.device)
+        )
+        x2_cropped = torch.index_select(
+            x2, dim=axis, index=torch.arange(dim_cropped).to(device=x1.device)
+        )
+        return x1_cropped, x2_cropped