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

braindecode/models/base.py

@@ -0,0 +1,296 @@
+# Authors: Pierre Guetschel
+#          Maciej Sliwowski
+#
+# License: BSD-3
+
+from __future__ import annotations
+
+import warnings
+from collections import OrderedDict
+from typing import Dict, Iterable, Optional
+
+import numpy as np
+import torch
+from docstring_inheritance import NumpyDocstringInheritanceInitMeta
+from torchinfo import ModelStatistics, summary
+
+
+def deprecated_args(obj, *old_new_args):
+    out_args = []
+    for old_name, new_name, old_val, new_val in old_new_args:
+        if old_val is None:
+            out_args.append(new_val)
+        else:
+            warnings.warn(
+                f"{obj.__class__.__name__}: {old_name!r} is depreciated. Use {new_name!r} instead."
+            )
+            if new_val is not None:
+                raise ValueError(
+                    f"{obj.__class__.__name__}: Both {old_name!r} and {new_name!r} were specified."
+                )
+            out_args.append(old_val)
+    return out_args
+
+
+class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
+    """
+    Mixin class for all EEG models in braindecode.
+
+    Parameters
+    ----------
+    n_outputs : int
+        Number of outputs of the model. This is the number of classes
+        in the case of classification.
+    n_chans : int
+        Number of EEG channels.
+    chs_info : list of dict
+        Information about each individual EEG channel. This should be filled with
+        ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+    n_times : int
+        Number of time samples of the input window.
+    input_window_seconds : float
+        Length of the input window in seconds.
+    sfreq : float
+        Sampling frequency of the EEG recordings.
+
+    Raises
+    ------
+    ValueError: If some input signal-related parameters are not specified
+                and can not be inferred.
+
+    Notes
+    -----
+    If some input signal-related parameters are not specified,
+    there will be an attempt to infer them from the other parameters.
+    """
+
+    def __init__(
+        self,
+        n_outputs: Optional[int] = None,  # type: ignore[assignment]
+        n_chans: Optional[int] = None,  # type: ignore[assignment]
+        chs_info=None,  # type: ignore[assignment]
+        n_times: Optional[int] = None,  # type: ignore[assignment]
+        input_window_seconds: Optional[float] = None,  # type: ignore[assignment]
+        sfreq: Optional[float] = None,  # type: ignore[assignment]
+    ):
+        if n_chans is not None and chs_info is not None and len(chs_info) != n_chans:
+            raise ValueError(f"{n_chans=} different from {chs_info=} length")
+        if (
+            n_times is not None
+            and input_window_seconds is not None
+            and sfreq is not None
+            and n_times != int(input_window_seconds * sfreq)
+        ):
+            raise ValueError(
+                f"{n_times=} different from {input_window_seconds=} * {sfreq=}"
+            )
+
+        self._input_window_seconds = input_window_seconds  # type: ignore[assignment]
+        self._chs_info = chs_info  # type: ignore[assignment]
+        self._n_outputs = n_outputs  # type: ignore[assignment]
+        self._n_chans = n_chans  # type: ignore[assignment]
+        self._n_times = n_times  # type: ignore[assignment]
+        self._sfreq = sfreq  # type: ignore[assignment]
+
+        super().__init__()
+
+    @property
+    def n_outputs(self) -> int:
+        if self._n_outputs is None:
+            raise ValueError("n_outputs not specified.")
+        return self._n_outputs
+
+    @property
+    def n_chans(self) -> int:
+        if self._n_chans is None and self._chs_info is not None:
+            return len(self._chs_info)
+        elif self._n_chans is None:
+            raise ValueError(
+                "n_chans could not be inferred. Either specify n_chans or chs_info."
+            )
+        return self._n_chans
+
+    @property
+    def chs_info(self) -> list[str]:
+        if self._chs_info is None:
+            raise ValueError("chs_info not specified.")
+        return self._chs_info
+
+    @property
+    def n_times(self) -> int:
+        if (
+            self._n_times is None
+            and self._input_window_seconds is not None
+            and self._sfreq is not None
+        ):
+            return int(self._input_window_seconds * self._sfreq)
+        elif self._n_times is None:
+            raise ValueError(
+                "n_times could not be inferred. "
+                "Either specify n_times or input_window_seconds and sfreq."
+            )
+        return self._n_times
+
+    @property
+    def input_window_seconds(self) -> float:
+        if (
+            self._input_window_seconds is None
+            and self._n_times is not None
+            and self._sfreq is not None
+        ):
+            return float(self._n_times / self._sfreq)
+        elif self._input_window_seconds is None:
+            raise ValueError(
+                "input_window_seconds could not be inferred. "
+                "Either specify input_window_seconds or n_times and sfreq."
+            )
+        return self._input_window_seconds
+
+    @property
+    def sfreq(self) -> float:
+        if (
+            self._sfreq is None
+            and self._input_window_seconds is not None
+            and self._n_times is not None
+        ):
+            return float(self._n_times / self._input_window_seconds)
+        elif self._sfreq is None:
+            raise ValueError(
+                "sfreq could not be inferred. "
+                "Either specify sfreq or input_window_seconds and n_times."
+            )
+        return self._sfreq
+
+    @property
+    def input_shape(self) -> tuple[int, int, int]:
+        """Input data shape."""
+        return (1, self.n_chans, self.n_times)
+
+    def get_output_shape(self) -> tuple[int, ...]:
+        """Returns shape of neural network output for batch size equal 1.
+
+        Returns
+        -------
+        output_shape: tuple[int, ...]
+            shape of the network output for `batch_size==1` (1, ...)
+        """
+        with torch.inference_mode():
+            try:
+                return tuple(
+                    self.forward(  # type: ignore
+                        torch.zeros(
+                            self.input_shape,
+                            dtype=next(self.parameters()).dtype,  # type: ignore
+                            device=next(self.parameters()).device,  # type: ignore
+                        )
+                    ).shape
+                )
+            except RuntimeError as exc:
+                if str(exc).endswith(
+                    (
+                        "Output size is too small",
+                        "Kernel size can't be greater than actual input size",
+                    )
+                ):
+                    msg = (
+                        "During model prediction RuntimeError was thrown showing that at some "
+                        f"layer `{str(exc).split('.')[-1]}` (see above in the stacktrace). This "
+                        "could be caused by providing too small `n_times`/`input_window_seconds`. "
+                        "Model may require longer chunks of signal in the input than "
+                        f"{self.input_shape}."
+                    )
+                    raise ValueError(msg) from exc
+                raise exc
+
+    mapping: Optional[Dict[str, str]] = None
+
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        mapping = self.mapping if self.mapping else {}
+        new_state_dict = OrderedDict()
+        for k, v in state_dict.items():
+            if k in mapping:
+                new_state_dict[mapping[k]] = v
+            else:
+                new_state_dict[k] = v
+
+        return super().load_state_dict(new_state_dict, *args, **kwargs)
+
+    def to_dense_prediction_model(self, axis: tuple[int, ...] | int = (2, 3)) -> None:
+        """
+        Transform a sequential model with strides to a model that outputs
+        dense predictions by removing the strides and instead inserting dilations.
+        Modifies model in-place.
+
+        Parameters
+        ----------
+        axis: int or (int,int)
+            Axis to transform (in terms of intermediate output axes)
+            can either be 2, 3, or (2,3).
+
+        Notes
+        -----
+        Does not yet work correctly for average pooling.
+        Prior to version 0.1.7, there had been a bug that could move strides
+        backwards one layer.
+
+        """
+        if not hasattr(axis, "__iter__"):
+            axis = (axis,)
+        assert all([ax in [2, 3] for ax in axis]), "Only 2 and 3 allowed for axis"  # type: ignore[union-attr]
+        axis = np.array(axis) - 2
+        stride_so_far = np.array([1, 1])
+        for module in self.modules():  # type: ignore
+            if hasattr(module, "dilation"):
+                assert module.dilation == 1 or (module.dilation == (1, 1)), (
+                    "Dilation should equal 1 before conversion, maybe the model is "
+                    "already converted?"
+                )
+                new_dilation = [1, 1]
+                for ax in axis:  # type: ignore[union-attr]
+                    new_dilation[ax] = int(stride_so_far[ax])
+                module.dilation = tuple(new_dilation)
+            if hasattr(module, "stride"):
+                if not hasattr(module.stride, "__len__"):
+                    module.stride = (module.stride, module.stride)
+                stride_so_far *= np.array(module.stride)
+                new_stride = list(module.stride)
+                for ax in axis:  # type: ignore[union-attr]
+                    new_stride[ax] = 1
+                module.stride = tuple(new_stride)
+
+    def get_torchinfo_statistics(
+        self,
+        col_names: Optional[Iterable[str]] = (
+            "input_size",
+            "output_size",
+            "num_params",
+            "kernel_size",
+        ),
+        row_settings: Optional[Iterable[str]] = ("var_names", "depth"),
+    ) -> ModelStatistics:
+        """Generate table describing the model using torchinfo.summary.
+
+        Parameters
+        ----------
+        col_names : tuple, optional
+            Specify which columns to show in the output, see torchinfo for details, by default
+            ("input_size", "output_size", "num_params", "kernel_size")
+        row_settings : tuple, optional
+             Specify which features to show in a row, see torchinfo for details, by default
+             ("var_names", "depth")
+
+        Returns
+        -------
+        torchinfo.ModelStatistics
+            ModelStatistics generated by torchinfo.summary.
+        """
+        return summary(
+            self,
+            input_size=(1, self.n_chans, self.n_times),
+            col_names=col_names,
+            row_settings=row_settings,
+            verbose=0,
+        )
+
+    def __str__(self) -> str:
+        return str(self.get_torchinfo_statistics())