braindecode 0.5.1__py3-none-any.whl

braindecode/models/deep4.py

@@ -0,0 +1,252 @@
+import numpy as np
+from torch import nn
+from torch.nn import init
+from torch.nn.functional import elu
+
+from .modules import Expression, AvgPool2dWithConv, Ensure4d
+from .functions import identity, transpose_time_to_spat, squeeze_final_output
+from ..util import np_to_var
+
+
+class Deep4Net(nn.Sequential):
+    """
+    Deep ConvNet model from [1]_.
+
+    References
+    ----------
+
+    .. [1] Schirrmeister, R. T., Springenberg, J. T., Fiederer, L. D. J.,
+       Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F. & Ball, T. (2017).
+       Deep learning with convolutional neural networks for EEG decoding and
+       visualization.
+       Human Brain Mapping , Aug. 2017. Online: http://dx.doi.org/10.1002/hbm.23730
+    """
+
+    def __init__(
+        self,
+        in_chans,
+        n_classes,
+        input_window_samples,
+        final_conv_length,
+        n_filters_time=25,
+        n_filters_spat=25,
+        filter_time_length=10,
+        pool_time_length=3,
+        pool_time_stride=3,
+        n_filters_2=50,
+        filter_length_2=10,
+        n_filters_3=100,
+        filter_length_3=10,
+        n_filters_4=200,
+        filter_length_4=10,
+        first_nonlin=elu,
+        first_pool_mode="max",
+        first_pool_nonlin=identity,
+        later_nonlin=elu,
+        later_pool_mode="max",
+        later_pool_nonlin=identity,
+        drop_prob=0.5,
+        double_time_convs=False,
+        split_first_layer=True,
+        batch_norm=True,
+        batch_norm_alpha=0.1,
+        stride_before_pool=False,
+    ):
+        super().__init__()
+        if final_conv_length == "auto":
+            assert input_window_samples is not None
+        self.in_chans = in_chans
+        self.n_classes = n_classes
+        self.input_window_samples = input_window_samples
+        self.final_conv_length = final_conv_length
+        self.n_filters_time = n_filters_time
+        self.n_filters_spat = n_filters_spat
+        self.filter_time_length = filter_time_length
+        self.pool_time_length = pool_time_length
+        self.pool_time_stride = pool_time_stride
+        self.n_filters_2 = n_filters_2
+        self.filter_length_2 = filter_length_2
+        self.n_filters_3 = n_filters_3
+        self.filter_length_3 = filter_length_3
+        self.n_filters_4 = n_filters_4
+        self.filter_length_4 = filter_length_4
+        self.first_nonlin = first_nonlin
+        self.first_pool_mode = first_pool_mode
+        self.first_pool_nonlin = first_pool_nonlin
+        self.later_nonlin = later_nonlin
+        self.later_pool_mode = later_pool_mode
+        self.later_pool_nonlin = later_pool_nonlin
+        self.drop_prob = drop_prob
+        self.double_time_convs = double_time_convs
+        self.split_first_layer = split_first_layer
+        self.batch_norm = batch_norm
+        self.batch_norm_alpha = batch_norm_alpha
+        self.stride_before_pool = stride_before_pool
+
+        if self.stride_before_pool:
+            conv_stride = self.pool_time_stride
+            pool_stride = 1
+        else:
+            conv_stride = 1
+            pool_stride = self.pool_time_stride
+        self.add_module("ensuredims", Ensure4d())
+        pool_class_dict = dict(max=nn.MaxPool2d, mean=AvgPool2dWithConv)
+        first_pool_class = pool_class_dict[self.first_pool_mode]
+        later_pool_class = pool_class_dict[self.later_pool_mode]
+        if self.split_first_layer:
+            self.add_module("dimshuffle", Expression(transpose_time_to_spat))
+            self.add_module(
+                "conv_time",
+                nn.Conv2d(
+                    1,
+                    self.n_filters_time,
+                    (self.filter_time_length, 1),
+                    stride=1,
+                ),
+            )
+            self.add_module(
+                "conv_spat",
+                nn.Conv2d(
+                    self.n_filters_time,
+                    self.n_filters_spat,
+                    (1, self.in_chans),
+                    stride=(conv_stride, 1),
+                    bias=not self.batch_norm,
+                ),
+            )
+            n_filters_conv = self.n_filters_spat
+        else:
+            self.add_module(
+                "conv_time",
+                nn.Conv2d(
+                    self.in_chans,
+                    self.n_filters_time,
+                    (self.filter_time_length, 1),
+                    stride=(conv_stride, 1),
+                    bias=not self.batch_norm,
+                ),
+            )
+            n_filters_conv = self.n_filters_time
+        if self.batch_norm:
+            self.add_module(
+                "bnorm",
+                nn.BatchNorm2d(
+                    n_filters_conv,
+                    momentum=self.batch_norm_alpha,
+                    affine=True,
+                    eps=1e-5,
+                ),
+            )
+        self.add_module("conv_nonlin", Expression(self.first_nonlin))
+        self.add_module(
+            "pool",
+            first_pool_class(
+                kernel_size=(self.pool_time_length, 1), stride=(pool_stride, 1)
+            ),
+        )
+        self.add_module("pool_nonlin", Expression(self.first_pool_nonlin))
+
+        def add_conv_pool_block(
+            model, n_filters_before, n_filters, filter_length, block_nr
+        ):
+            suffix = "_{:d}".format(block_nr)
+            self.add_module("drop" + suffix, nn.Dropout(p=self.drop_prob))
+            self.add_module(
+                "conv" + suffix,
+                nn.Conv2d(
+                    n_filters_before,
+                    n_filters,
+                    (filter_length, 1),
+                    stride=(conv_stride, 1),
+                    bias=not self.batch_norm,
+                ),
+            )
+            if self.batch_norm:
+                self.add_module(
+                    "bnorm" + suffix,
+                    nn.BatchNorm2d(
+                        n_filters,
+                        momentum=self.batch_norm_alpha,
+                        affine=True,
+                        eps=1e-5,
+                    ),
+                )
+            self.add_module("nonlin" + suffix, Expression(self.later_nonlin))
+
+            self.add_module(
+                "pool" + suffix,
+                later_pool_class(
+                    kernel_size=(self.pool_time_length, 1),
+                    stride=(pool_stride, 1),
+                ),
+            )
+            self.add_module(
+                "pool_nonlin" + suffix, Expression(self.later_pool_nonlin)
+            )
+
+        add_conv_pool_block(
+            self, n_filters_conv, self.n_filters_2, self.filter_length_2, 2
+        )
+        add_conv_pool_block(
+            self, self.n_filters_2, self.n_filters_3, self.filter_length_3, 3
+        )
+        add_conv_pool_block(
+            self, self.n_filters_3, self.n_filters_4, self.filter_length_4, 4
+        )
+
+        # self.add_module('drop_classifier', nn.Dropout(p=self.drop_prob))
+        self.eval()
+        if self.final_conv_length == "auto":
+            out = self(
+                np_to_var(
+                    np.ones(
+                        (1, self.in_chans, self.input_window_samples, 1),
+                        dtype=np.float32,
+                    )
+                )
+            )
+            n_out_time = out.cpu().data.numpy().shape[2]
+            self.final_conv_length = n_out_time
+        self.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                self.n_filters_4,
+                self.n_classes,
+                (self.final_conv_length, 1),
+                bias=True,
+            ),
+        )
+        self.add_module("softmax", nn.LogSoftmax(dim=1))
+        self.add_module("squeeze", Expression(squeeze_final_output))
+
+        # Initialization, xavier is same as in our paper...
+        # was default from lasagne
+        init.xavier_uniform_(self.conv_time.weight, gain=1)
+        # maybe no bias in case of no split layer and batch norm
+        if self.split_first_layer or (not self.batch_norm):
+            init.constant_(self.conv_time.bias, 0)
+        if self.split_first_layer:
+            init.xavier_uniform_(self.conv_spat.weight, gain=1)
+            if not self.batch_norm:
+                init.constant_(self.conv_spat.bias, 0)
+        if self.batch_norm:
+            init.constant_(self.bnorm.weight, 1)
+            init.constant_(self.bnorm.bias, 0)
+        param_dict = dict(list(self.named_parameters()))
+        for block_nr in range(2, 5):
+            conv_weight = param_dict["conv_{:d}.weight".format(block_nr)]
+            init.xavier_uniform_(conv_weight, gain=1)
+            if not self.batch_norm:
+                conv_bias = param_dict["conv_{:d}.bias".format(block_nr)]
+                init.constant_(conv_bias, 0)
+            else:
+                bnorm_weight = param_dict["bnorm_{:d}.weight".format(block_nr)]
+                bnorm_bias = param_dict["bnorm_{:d}.bias".format(block_nr)]
+                init.constant_(bnorm_weight, 1)
+                init.constant_(bnorm_bias, 0)
+
+        init.xavier_uniform_(self.conv_classifier.weight, gain=1)
+        init.constant_(self.conv_classifier.bias, 0)
+
+        # Start in eval mode
+        self.eval()

braindecode/models/eegnet.py

@@ -0,0 +1,339 @@
+import torch
+from torch import nn
+from torch.nn.functional import elu
+
+from .modules import Expression, Ensure4d
+from .functions import squeeze_final_output
+
+class Conv2dWithConstraint(nn.Conv2d):
+    def __init__(self, *args, max_norm=1, **kwargs):
+        self.max_norm = max_norm
+        super(Conv2dWithConstraint, self).__init__(*args, **kwargs)
+
+    def forward(self, x):
+        self.weight.data = torch.renorm(
+            self.weight.data, p=2, dim=0, maxnorm=self.max_norm
+        )
+        return super(Conv2dWithConstraint, self).forward(x)
+
+
+class EEGNetv4(nn.Sequential):
+    """
+    EEGNet v4 model from [EEGNet4]_.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper description.
+
+    References
+    ----------
+
+    .. [EEGNet4] Lawhern, V. J., Solon, A. J., Waytowich, N. R., Gordon,
+       S. M., Hung, C. P., & Lance, B. J. (2018).
+       EEGNet: A Compact Convolutional Network for EEG-based
+       Brain-Computer Interfaces.
+       arXiv preprint arXiv:1611.08024.
+    """
+
+    def __init__(
+        self,
+        in_chans,
+        n_classes,
+        input_window_samples=None,
+        final_conv_length="auto",
+        pool_mode="mean",
+        F1=8,
+        D=2,
+        F2=16,  # usually set to F1*D (?)
+        kernel_length=64,
+        third_kernel_size=(8, 4),
+        drop_prob=0.25,
+    ):
+        super().__init__()
+        if final_conv_length == "auto":
+            assert input_window_samples is not None
+        self.in_chans = in_chans
+        self.n_classes = n_classes
+        self.input_window_samples = input_window_samples
+        self.final_conv_length = final_conv_length
+        self.pool_mode = pool_mode
+        self.F1 = F1
+        self.D = D
+        self.F2 = F2
+        self.kernel_length = kernel_length
+        self.third_kernel_size = third_kernel_size
+        self.drop_prob = drop_prob
+
+        pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.pool_mode]
+        self.add_module("ensuredims", Ensure4d())
+        # b c 0 1
+        # now to b 1 0 c
+        self.add_module("dimshuffle", Expression(_transpose_to_b_1_c_0))
+
+        self.add_module(
+            "conv_temporal",
+            nn.Conv2d(
+                1,
+                self.F1,
+                (1, self.kernel_length),
+                stride=1,
+                bias=False,
+                padding=(0, self.kernel_length // 2),
+            ),
+        )
+        self.add_module(
+            "bnorm_temporal",
+            nn.BatchNorm2d(self.F1, momentum=0.01, affine=True, eps=1e-3),
+        )
+        self.add_module(
+            "conv_spatial",
+            Conv2dWithConstraint(
+                self.F1,
+                self.F1 * self.D,
+                (self.in_chans, 1),
+                max_norm=1,
+                stride=1,
+                bias=False,
+                groups=self.F1,
+                padding=(0, 0),
+            ),
+        )
+
+        self.add_module(
+            "bnorm_1",
+            nn.BatchNorm2d(
+                self.F1 * self.D, momentum=0.01, affine=True, eps=1e-3
+            ),
+        )
+        self.add_module("elu_1", Expression(elu))
+
+        self.add_module("pool_1", pool_class(kernel_size=(1, 4), stride=(1, 4)))
+        self.add_module("drop_1", nn.Dropout(p=self.drop_prob))
+
+        # https://discuss.pytorch.org/t/how-to-modify-a-conv2d-to-depthwise-separable-convolution/15843/7
+        self.add_module(
+            "conv_separable_depth",
+            nn.Conv2d(
+                self.F1 * self.D,
+                self.F1 * self.D,
+                (1, 16),
+                stride=1,
+                bias=False,
+                groups=self.F1 * self.D,
+                padding=(0, 16 // 2),
+            ),
+        )
+        self.add_module(
+            "conv_separable_point",
+            nn.Conv2d(
+                self.F1 * self.D,
+                self.F2,
+                (1, 1),
+                stride=1,
+                bias=False,
+                padding=(0, 0),
+            ),
+        )
+
+        self.add_module(
+            "bnorm_2",
+            nn.BatchNorm2d(self.F2, momentum=0.01, affine=True, eps=1e-3),
+        )
+        self.add_module("elu_2", Expression(elu))
+        self.add_module("pool_2", pool_class(kernel_size=(1, 8), stride=(1, 8)))
+        self.add_module("drop_2", nn.Dropout(p=self.drop_prob))
+
+        out = self(
+            torch.ones(
+                (1, self.in_chans, self.input_window_samples, 1),
+                dtype=torch.float32
+            )
+        )
+        n_out_virtual_chans = out.cpu().data.numpy().shape[2]
+
+        if self.final_conv_length == "auto":
+            n_out_time = out.cpu().data.numpy().shape[3]
+            self.final_conv_length = n_out_time
+
+        self.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                self.F2,
+                self.n_classes,
+                (n_out_virtual_chans, self.final_conv_length),
+                bias=True,
+            ),
+        )
+        self.add_module("softmax", nn.LogSoftmax(dim=1))
+        # Transpose back to the the logic of braindecode,
+        # so time in third dimension (axis=2)
+        self.add_module("permute_back", Expression(_transpose_1_0))
+        self.add_module("squeeze", Expression(squeeze_final_output))
+
+        _glorot_weight_zero_bias(self)
+
+
+def _transpose_to_b_1_c_0(x):
+    return x.permute(0, 3, 1, 2)
+
+
+def _transpose_1_0(x):
+    return x.permute(0, 1, 3, 2)
+
+
+class EEGNetv1(nn.Sequential):
+    """
+    EEGNet model from [EEGNet]_.
+
+    Notes
+    -----
+    This implementation is not guaranteed to be correct, has not been checked
+    by original authors, only reimplemented from the paper description.
+
+    References
+    ----------
+
+    .. [EEGNet] Lawhern, V. J., Solon, A. J., Waytowich, N. R., Gordon,
+       S. M., Hung, C. P., & Lance, B. J. (2016).
+       EEGNet: A Compact Convolutional Network for EEG-based
+       Brain-Computer Interfaces.
+       arXiv preprint arXiv:1611.08024.
+    """
+
+    def __init__(
+        self,
+        in_chans,
+        n_classes,
+        input_window_samples=None,
+        final_conv_length="auto",
+        pool_mode="max",
+        second_kernel_size=(2, 32),
+        third_kernel_size=(8, 4),
+        drop_prob=0.25,
+    ):
+        super().__init__()
+        if final_conv_length == "auto":
+            assert input_window_samples is not None
+        self.in_chans = in_chans
+        self.n_classes = n_classes
+        self.input_window_samples = input_window_samples
+        self.final_conv_length = final_conv_length
+        self.pool_mode = pool_mode
+        self.second_kernel_size = second_kernel_size
+        self.third_kernel_size = third_kernel_size
+        self.drop_prob = drop_prob
+
+        pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.pool_mode]
+        self.add_module("ensuredims", Ensure4d())
+        n_filters_1 = 16
+        self.add_module(
+            "conv_1",
+            nn.Conv2d(self.in_chans, n_filters_1, (1, 1), stride=1, bias=True),
+        )
+        self.add_module(
+            "bnorm_1",
+            nn.BatchNorm2d(n_filters_1, momentum=0.01, affine=True, eps=1e-3),
+        )
+        self.add_module("elu_1", Expression(elu))
+        # transpose to examples x 1 x (virtual, not EEG) channels x time
+        self.add_module(
+            "permute_1", Expression(lambda x: x.permute(0, 3, 1, 2))
+        )
+
+        self.add_module("drop_1", nn.Dropout(p=self.drop_prob))
+
+        n_filters_2 = 4
+        # keras padds unequal padding more in front, so padding
+        # too large should be ok.
+        # Not padding in time so that cropped training makes sense
+        # https://stackoverflow.com/questions/43994604/padding-with-even-kernel-size-in-a-convolutional-layer-in-keras-theano
+
+        self.add_module(
+            "conv_2",
+            nn.Conv2d(
+                1,
+                n_filters_2,
+                self.second_kernel_size,
+                stride=1,
+                padding=(self.second_kernel_size[0] // 2, 0),
+                bias=True,
+            ),
+        )
+        self.add_module(
+            "bnorm_2",
+            nn.BatchNorm2d(n_filters_2, momentum=0.01, affine=True, eps=1e-3),
+        )
+        self.add_module("elu_2", Expression(elu))
+        self.add_module("pool_2", pool_class(kernel_size=(2, 4), stride=(2, 4)))
+        self.add_module("drop_2", nn.Dropout(p=self.drop_prob))
+
+        n_filters_3 = 4
+        self.add_module(
+            "conv_3",
+            nn.Conv2d(
+                n_filters_2,
+                n_filters_3,
+                self.third_kernel_size,
+                stride=1,
+                padding=(self.third_kernel_size[0] // 2, 0),
+                bias=True,
+            ),
+        )
+        self.add_module(
+            "bnorm_3",
+            nn.BatchNorm2d(n_filters_3, momentum=0.01, affine=True, eps=1e-3),
+        )
+        self.add_module("elu_3", Expression(elu))
+        self.add_module("pool_3", pool_class(kernel_size=(2, 4), stride=(2, 4)))
+        self.add_module("drop_3", nn.Dropout(p=self.drop_prob))
+
+        out = self(
+            torch.ones(
+                    (1, self.in_chans, self.input_window_samples, 1),
+                    dtype=torch.float32,
+            )
+        )
+        n_out_virtual_chans = out.cpu().data.numpy().shape[2]
+
+        if self.final_conv_length == "auto":
+            n_out_time = out.cpu().data.numpy().shape[3]
+            self.final_conv_length = n_out_time
+
+        self.add_module(
+            "conv_classifier",
+            nn.Conv2d(
+                n_filters_3,
+                self.n_classes,
+                (n_out_virtual_chans, self.final_conv_length),
+                bias=True,
+            ),
+        )
+        self.add_module("softmax", nn.LogSoftmax(dim=1))
+        # Transpose back to the the logic of braindecode,
+        # so time in third dimension (axis=2)
+        self.add_module(
+            "permute_2", Expression(lambda x: x.permute(0, 1, 3, 2))
+        )
+        self.add_module("squeeze", Expression(squeeze_final_output))
+        _glorot_weight_zero_bias(self)
+
+def _glorot_weight_zero_bias(model):
+    """Initalize parameters of all modules by initializing weights with
+    glorot
+     uniform/xavier initialization, and setting biases to zero. Weights from
+     batch norm layers are set to 1.
+
+    Parameters
+    ----------
+    model: Module
+    """
+    for module in model.modules():
+        if hasattr(module, "weight"):
+            if not ("BatchNorm" in module.__class__.__name__):
+                nn.init.xavier_uniform_(module.weight, gain=1)
+            else:
+                nn.init.constant_(module.weight, 1)
+        if hasattr(module, "bias"):
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)