Adding Braindecode pipeline

moabb/pipelines/__init__.py

@@ -3,6 +3,8 @@
 Pipelines are typically a chain of sklearn compatible transformers and end
 with a sklearn compatible estimator.
 """
+from .braindecode import BraindecodeClassifierModel, CreateBraindecodeDataset
+
 # flake8: noqa
 from .classification import SSVEP_CCA, SSVEP_TRCA
 from .features import FM, ExtendedSSVEPSignal, LogVariance

moabb/pipelines/braindecode.py

@@ -0,0 +1,76 @@
+from braindecode.datasets import WindowsDataset, create_from_X_y
+from mne.epochs import BaseEpochs
+from numpy import array, unique
+from sklearn.base import BaseEstimator, ClassifierMixin, TransformerMixin
+
+
+class CreateBraindecodeDataset(BaseEstimator, TransformerMixin):
+    """
+    Wrapper to create a Braindecode Dataset from a mne Epoched
+    object.
+
+    This is a transformer function that allow used to use the
+    dataset as a sklearn pipeline.
+    """
+
+    def __init__(self, kw_args: dict = None):
+        """
+
+        Parameters
+        ----------
+        kw_args: dict
+
+        """
+        self.kw_args = kw_args
+
+    def fit(self, X: BaseEpochs, y=None):
+        self.y = y
+        return self
+
+    def transform(self, X: BaseEpochs, y=None) -> WindowsDataset:
+        """
+
+        Parameters
+        ----------
+        X: BaseEpochs object from mne
+        y: list|array of labels
+
+        Returns
+        -------
+        WindowsDataset: Braindecode Dataset
+        """
+        dataset = create_from_X_y(
+            X.get_data(),
+            y=self.y,
+            window_size_samples=X.get_data().shape[2],
+            window_stride_samples=X.get_data().shape[2],
+            drop_last_window=False,
+            sfreq=X.info["sfreq"],
+        )
+
+        return dataset
+
+    def __sklearn_is_fitted__(self) -> bool:
+        """
+        Return True since CreateBraindecodeDataset is stateless.
+        """
+        return True
+
+
+class BraindecodeClassifierModel(BaseEstimator, ClassifierMixin):
+    def __init__(self, clf: BaseEstimator, kw_args: dict = None):
+        self.clf = clf
+        self.classes_ = None
+        self.kw_args = kw_args
+
+    def fit(self, X: WindowsDataset, y=None) -> BaseEstimator:
+        self.clf.fit(X, y=y, **self.kw_args)
+        self.classes_ = unique(y)
+
+        return self.clf
+
+    def predict(self, X: WindowsDataset) -> array:
+        return self.clf.predict(X)
+
+    def predict_proba(self, X: WindowsDataset) -> array:
+        return self.clf.predict_proba(X)

pipelines/braindecode.py

@@ -0,0 +1,62 @@
+from braindecode.classifier import EEGClassifier
+from braindecode.models import ShallowFBCSPNet
+from sklearn.pipeline import Pipeline
+from skorch.callbacks import LRScheduler
+from skorch.dataset import ValidSplit
+from torch.cuda import is_available
+from torch.nn import NLLLoss
+from torch.optim import AdamW
+
+from moabb.pipelines.braindecode import (
+    BraindecodeClassifierModel,
+    CreateBraindecodeDataset,
+)
+
+
+# hard-coded for now
+n_classes = 2
+n_chans = 22
+input_window_samples = 1001
+# These values we found good for shallow network:
+lr = 0.0625 * 0.01
+weight_decay = 0
+batch_size = 64
+n_epochs = 4
+
+model = ShallowFBCSPNet(
+    n_chans,
+    n_classes,
+    input_window_samples=input_window_samples,
+    final_conv_length="auto",
+)
+device = "cuda" if is_available() else "cpu"
+
+clf = EEGClassifier(
+    model,
+    criterion=NLLLoss,
+    optimizer=AdamW,
+    train_split=ValidSplit(0.2),  # using valid_set for validation
+    optimizer__lr=lr,
+    optimizer__weight_decay=weight_decay,
+    batch_size=batch_size,
+    callbacks=[
+        "accuracy",
+        ("lr_scheduler", LRScheduler("CosineAnnealingLR", T_max=n_epochs - 1)),
+    ],
+    device=device,
+)
+
+create_dataset = CreateBraindecodeDataset()
+fit_params = {"epochs": 10}
+
+clf_braindecode = BraindecodeClassifierModel(clf, fit_params)
+
+pipe = Pipeline([("Braindecode_dataset", create_dataset), ("Net", clf_braindecode)])
+
+pipes = {"ShallowFBCSPNet": pipe}
+
+PIPELINE = {
+    "name": "ShallowFBCSPNet",
+    "paradigms": ["LeftRightImagery"],
+    "pipeline": pipe,
+}

pyproject.toml

@@ -49,6 +49,8 @@ pre-commit = "^2.21.0"
 m2r2 = "^0.3.3"
 tdlda = {git = "https://github.com/jsosulski/tdlda.git", rev = "0.1.0"}
 
+
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"