braindecode 0.5.1__py3-none-any.whl

braindecode/models/eegresnet.py

@@ -0,0 +1,232 @@
+# Authors: Robin Tibor Schirrmeister
+#          Tonio Ball
+#
+# License: BSD-3
+
+import numpy as np
+
+import torch
+from torch import nn
+from torch.nn import init
+from torch.nn.functional import elu
+
+from .functions import transpose_time_to_spat, squeeze_final_output
+from ..util import np_to_var
+from .modules import Expression, AvgPool2dWithConv, Ensure4d
+
+
+class EEGResNet(nn.Sequential):
+    """
+    Residual Network for EEG.
+    """
+    def __init__(self,
+                 in_chans,
+                 n_classes,
+                 input_window_samples,
+                 final_pool_length,
+                 n_first_filters,
+                 n_layers_per_block=2,
+                 first_filter_length=3,
+                 nonlinearity=elu,
+                 split_first_layer=True,
+                 batch_norm_alpha=0.1,
+                 batch_norm_epsilon=1e-4,
+                 conv_weight_init_fn=lambda w: init.kaiming_normal_(w, a=0)):
+        super().__init__()
+        self.in_chans = in_chans
+        self.n_classes = n_classes
+        self.input_window_samples = input_window_samples
+        if final_pool_length == 'auto':
+            assert input_window_samples is not None
+        assert first_filter_length % 2 == 1
+        self.final_pool_length = final_pool_length
+        self.n_first_filters = n_first_filters
+        self.n_layers_per_block = n_layers_per_block
+        self.first_filter_length = first_filter_length
+        self.nonlinearity = nonlinearity
+        self.split_first_layer = split_first_layer
+        self.batch_norm_alpha = batch_norm_alpha
+        self.batch_norm_epsilon = batch_norm_epsilon
+        self.conv_weight_init_fn = conv_weight_init_fn
+
+        self.add_module("ensuredims", Ensure4d())
+        if self.split_first_layer:
+            self.add_module('dimshuffle', Expression(transpose_time_to_spat))
+            self.add_module('conv_time', nn.Conv2d(1, self.n_first_filters,
+                                                   (self.first_filter_length, 1),
+                                                   stride=1,
+                                                   padding=(self.first_filter_length // 2, 0)))
+            self.add_module('conv_spat',
+                            nn.Conv2d(self.n_first_filters, self.n_first_filters,
+                                      (1, self.in_chans),
+                                      stride=(1, 1),
+                                      bias=False))
+        else:
+            self.add_module('conv_time',
+                            nn.Conv2d(self.in_chans, self.n_first_filters,
+                                      (self.first_filter_length, 1),
+                                      stride=(1, 1),
+                                      padding=(self.first_filter_length // 2, 0),
+                                      bias=False,))
+        n_filters_conv = self.n_first_filters
+        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.nonlinearity))
+        cur_dilation = np.array([1, 1])
+        n_cur_filters = n_filters_conv
+        i_block = 1
+        for i_layer in range(self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+        i_block += 1
+        cur_dilation[0] *= 2
+        n_out_filters = int(2 * n_cur_filters)
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_out_filters,
+                                       dilation=cur_dilation,))
+        n_cur_filters = n_out_filters
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        n_out_filters = int(1.5 * n_cur_filters)
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_out_filters,
+                                       dilation=cur_dilation,))
+        n_cur_filters = n_out_filters
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_cur_filters,
+                                       dilation=cur_dilation,))
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_cur_filters,
+                                       dilation=cur_dilation,))
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_cur_filters,
+                                       dilation=cur_dilation,))
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+        i_block += 1
+        cur_dilation[0] *= 2
+        self.add_module('res_{:d}_{:d}'.format(i_block, 0),
+                        _ResidualBlock(n_cur_filters, n_cur_filters,
+                                       dilation=cur_dilation,))
+        for i_layer in range(1, self.n_layers_per_block):
+            self.add_module('res_{:d}_{:d}'.format(i_block, i_layer),
+                            _ResidualBlock(n_cur_filters, n_cur_filters,
+                                           dilation=cur_dilation))
+
+        self.eval()
+        if self.final_pool_length == 'auto':
+            self.add_module('mean_pool', nn.AdaptiveAvgPool2d((1, 1)))
+        else:
+            pool_dilation = int(cur_dilation[0]), int(cur_dilation[1])
+            self.add_module('mean_pool', AvgPool2dWithConv(
+                (self.final_pool_length, 1), (1, 1),
+                dilation=pool_dilation))
+        self.add_module('conv_classifier',
+                        nn.Conv2d(n_cur_filters, self.n_classes,
+                                  (1, 1), bias=True))
+        self.add_module('softmax', nn.LogSoftmax(dim=1))
+        self.add_module('squeeze', Expression(squeeze_final_output))
+
+        # Initialize all weights
+        self.apply(lambda module: _weights_init(module, self.conv_weight_init_fn))
+
+        # Start in eval mode
+        self.eval()
+
+
+def _weights_init(module, conv_weight_init_fn):
+    """
+    initialize weights
+    """
+    classname = module.__class__.__name__
+    if 'Conv' in classname and classname != "AvgPool2dWithConv":
+        conv_weight_init_fn(module.weight)
+        if module.bias is not None:
+            init.constant_(module.bias, 0)
+    elif 'BatchNorm' in classname:
+        init.constant_(module.weight, 1)
+        init.constant_(module.bias, 0)
+
+
+class _ResidualBlock(nn.Module):
+    """
+    create a residual learning building block with two stacked 3x3 convlayers as in paper
+    """
+    def __init__(self, in_filters,
+                 out_num_filters,
+                 dilation,
+                 filter_time_length=3,
+                 nonlinearity=elu,
+                 batch_norm_alpha=0.1, batch_norm_epsilon=1e-4):
+        super(_ResidualBlock, self).__init__()
+        time_padding = int((filter_time_length - 1) * dilation[0])
+        assert time_padding % 2 == 0
+        time_padding = int(time_padding // 2)
+        dilation = (int(dilation[0]), int(dilation[1]))
+        assert (out_num_filters - in_filters) % 2 == 0, (
+            "Need even number of extra channels in order to be able to "
+            "pad correctly")
+        self.n_pad_chans = out_num_filters - in_filters
+
+        self.conv_1 = nn.Conv2d(
+            in_filters, out_num_filters, (filter_time_length, 1), stride=(1, 1),
+            dilation=dilation,
+            padding=(time_padding, 0))
+        self.bn1 = nn.BatchNorm2d(
+            out_num_filters, momentum=batch_norm_alpha, affine=True,
+            eps=batch_norm_epsilon)
+        self.conv_2 = nn.Conv2d(
+            out_num_filters, out_num_filters, (filter_time_length, 1), stride=(1, 1),
+            dilation=dilation,
+            padding=(time_padding, 0))
+        self.bn2 = nn.BatchNorm2d(
+            out_num_filters, momentum=batch_norm_alpha,
+            affine=True, eps=batch_norm_epsilon)
+        # also see https://mail.google.com/mail/u/0/#search/ilya+joos/1576137dd34c3127
+        # for resnet options as ilya used them
+        self.nonlinearity = nonlinearity
+
+    def forward(self, x):
+        stack_1 = self.nonlinearity(self.bn1(self.conv_1(x)))
+        stack_2 = self.bn2(self.conv_2(stack_1))  # next nonlin after sum
+        if self.n_pad_chans != 0:
+            zeros_for_padding = torch.autograd.Variable(
+                torch.zeros(x.size()[0], self.n_pad_chans // 2,
+                         x.size()[2], x.size()[3]))
+            if x.is_cuda:
+                zeros_for_padding = zeros_for_padding.cuda()
+            x = torch.cat((zeros_for_padding, x, zeros_for_padding), dim=1)
+        out = self.nonlinearity(x + stack_2)
+        return out

braindecode/models/functions.py

@@ -0,0 +1,43 @@
+import torch
+
+
+def square(x):
+    return x * x
+
+
+def safe_log(x, eps=1e-6):
+    """ Prevents :math:`log(0)` by using :math:`log(max(x, eps))`."""
+    return torch.log(torch.clamp(x, min=eps))
+
+
+def identity(x):
+    return x
+
+
+def squeeze_final_output(x):
+    """Removes empty dimension at end and potentially removes empty time
+     dimension. It does  not just use squeeze as we never want to remove
+     first dimension.
+
+    Returns
+    -------
+    x: torch.Tensor
+        squeezed tensor
+    """
+
+    assert x.size()[3] == 1
+    x = x[:, :, :, 0]
+    if x.size()[2] == 1:
+        x = x[:, :, 0]
+    return x
+
+
+def transpose_time_to_spat(x):
+    """Swap time and spatial dimensions.
+
+    Returns
+    -------
+    x: torch.Tensor
+        tensor in which last and first dimensions are swapped
+    """
+    return x.permute(0, 3, 2, 1)

braindecode/models/hybrid.py

@@ -0,0 +1,103 @@
+import torch
+from torch import nn
+from torch.nn import ConstantPad2d
+
+from .deep4 import Deep4Net
+from .util import to_dense_prediction_model
+from .shallow_fbcsp import ShallowFBCSPNet
+
+
+class HybridNet(nn.Module):
+    """Hybrid ConvNet model from [3]_.
+
+    Very hardcoded at the moment.
+
+    References
+    ----------
+
+    .. [3] 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):
+        super(HybridNet, self).__init__()
+        deep_model = Deep4Net(
+            in_chans,
+            n_classes,
+            n_filters_time=20,
+            n_filters_spat=30,
+            n_filters_2=40,
+            n_filters_3=50,
+            n_filters_4=60,
+            input_window_samples=input_window_samples,
+            final_conv_length=2,
+        )
+        shallow_model = ShallowFBCSPNet(
+            in_chans,
+            n_classes,
+            input_window_samples=input_window_samples,
+            n_filters_time=30,
+            n_filters_spat=40,
+            filter_time_length=28,
+            final_conv_length=29,
+        )
+
+        reduced_deep_model = nn.Sequential()
+        for name, module in deep_model.named_children():
+            if name == "conv_classifier":
+                new_conv_layer = nn.Conv2d(
+                    module.in_channels,
+                    60,
+                    kernel_size=module.kernel_size,
+                    stride=module.stride,
+                )
+                reduced_deep_model.add_module("deep_final_conv", new_conv_layer)
+                break
+            reduced_deep_model.add_module(name, module)
+
+        reduced_shallow_model = nn.Sequential()
+        for name, module in shallow_model.named_children():
+            if name == "conv_classifier":
+                new_conv_layer = nn.Conv2d(
+                    module.in_channels,
+                    40,
+                    kernel_size=module.kernel_size,
+                    stride=module.stride,
+                )
+                reduced_shallow_model.add_module(
+                    "shallow_final_conv", new_conv_layer
+                )
+                break
+            reduced_shallow_model.add_module(name, module)
+
+        to_dense_prediction_model(reduced_deep_model)
+        to_dense_prediction_model(reduced_shallow_model)
+        self.reduced_deep_model = reduced_deep_model
+        self.reduced_shallow_model = reduced_shallow_model
+        self.final_conv = nn.Conv2d(
+            100, n_classes, kernel_size=(1, 1), stride=1
+        )
+
+    def forward(self, x):
+        deep_out = self.reduced_deep_model(x)
+        shallow_out = self.reduced_shallow_model(x)
+
+        n_diff_deep_shallow = deep_out.size()[2] - shallow_out.size()[2]
+
+        if n_diff_deep_shallow < 0:
+            deep_out = ConstantPad2d((0, 0, -n_diff_deep_shallow, 0), 0)(
+                deep_out
+            )
+        elif n_diff_deep_shallow > 0:
+            shallow_out = ConstantPad2d((0, 0, n_diff_deep_shallow, 0), 0)(
+                shallow_out
+            )
+
+        merged_out = torch.cat((deep_out, shallow_out), dim=1)
+        linear_out = self.final_conv(merged_out)
+        softmaxed = nn.LogSoftmax(dim=1)(linear_out)
+        squeezed = softmaxed.squeeze(3)
+        return squeezed

braindecode/models/modules.py

@@ -0,0 +1,150 @@
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from ..util import np_to_var
+
+
+class Ensure4d(torch.nn.Module):
+    def forward(self, x):
+        while(len(x.shape) < 4):
+            x = x.unsqueeze(-1)
+        return x
+
+
+class Expression(torch.nn.Module):
+    """Compute given expression on forward pass.
+
+    Parameters
+    ----------
+    expression_fn : callable
+        Should accept variable number of objects of type
+        `torch.autograd.Variable` to compute its output.
+    """
+
+    def __init__(self, expression_fn):
+        super(Expression, self).__init__()
+        self.expression_fn = expression_fn
+
+    def forward(self, *x):
+        return self.expression_fn(*x)
+
+    def __repr__(self):
+        if hasattr(self.expression_fn, "func") and hasattr(
+            self.expression_fn, "kwargs"
+        ):
+            expression_str = "{:s} {:s}".format(
+                self.expression_fn.func.__name__, str(self.expression_fn.kwargs)
+            )
+        elif hasattr(self.expression_fn, "__name__"):
+            expression_str = self.expression_fn.__name__
+        else:
+            expression_str = repr(self.expression_fn)
+        return (
+            self.__class__.__name__ +
+            "(expression=%s) " % expression_str
+        )
+
+
+class AvgPool2dWithConv(torch.nn.Module):
+    """
+    Compute average pooling using a convolution, to have the dilation parameter.
+
+    Parameters
+    ----------
+    kernel_size: (int,int)
+        Size of the pooling region.
+    stride: (int,int)
+        Stride of the pooling operation.
+    dilation: int or (int,int)
+        Dilation applied to the pooling filter.
+    padding: int or (int,int)
+        Padding applied before the pooling operation.
+    """
+
+    def __init__(self, kernel_size, stride, dilation=1, padding=0):
+        super(AvgPool2dWithConv, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.dilation = dilation
+        self.padding = padding
+        # don't name them "weights" to
+        # make sure these are not accidentally used by some procedure
+        # that initializes parameters or something
+        self._pool_weights = None
+
+    def forward(self, x):
+        # Create weights for the convolution on demand:
+        # size or type of x changed...
+        in_channels = x.size()[1]
+        weight_shape = (
+            in_channels,
+            1,
+            self.kernel_size[0],
+            self.kernel_size[1],
+        )
+        if self._pool_weights is None or (
+            (tuple(self._pool_weights.size()) != tuple(weight_shape)) or
+            (self._pool_weights.is_cuda != x.is_cuda) or
+            (self._pool_weights.data.type() != x.data.type())
+        ):
+            n_pool = np.prod(self.kernel_size)
+            weights = np_to_var(
+                np.ones(weight_shape, dtype=np.float32) / float(n_pool)
+            )
+            weights = weights.type_as(x)
+            if x.is_cuda:
+                weights = weights.cuda()
+            self._pool_weights = weights
+
+        pooled = F.conv2d(
+            x,
+            self._pool_weights,
+            bias=None,
+            stride=self.stride,
+            dilation=self.dilation,
+            padding=self.padding,
+            groups=in_channels,
+        )
+        return pooled
+
+
+class IntermediateOutputWrapper(torch.nn.Module):
+    """Wraps network model such that outputs of intermediate layers can be returned.
+    forward() returns list of intermediate activations in a network during forward pass.
+
+    Parameters
+    ----------
+    to_select : list
+        list of module names for which activation should be returned
+    model : model object
+        network model
+
+    Examples
+    --------
+    >>> model = Deep4Net()
+    >>> select_modules = ['conv_spat','conv_2','conv_3','conv_4'] # Specify intermediate outputs
+    >>> model_pert = IntermediateOutputWrapper(select_modules,model) # Wrap model
+    """
+
+    def __init__(self, to_select, model):
+        if not len(list(model.children())) == len(list(model.named_children())):
+            raise Exception("All modules in model need to have names!")
+
+        super(IntermediateOutputWrapper, self).__init__()
+
+        modules_list = model.named_children()
+        for key, module in modules_list:
+            self.add_module(key, module)
+            self._modules[key].load_state_dict(module.state_dict())
+        self._to_select = to_select
+
+    def forward(self, x):
+        # Call modules individually and append activation to output if module is in to_select
+        o = []
+        for name, module in self._modules.items():
+            x = module(x)
+            if name in self._to_select:
+                o.append(x)
+        return o

braindecode/models/shallow_fbcsp.py

@@ -0,0 +1,158 @@
+import numpy as np
+from torch import nn
+from torch.nn import init
+
+from ..util import np_to_var
+from .modules import Expression, Ensure4d
+from .functions import (
+    safe_log, square, transpose_time_to_spat, squeeze_final_output
+)
+
+
+class ShallowFBCSPNet(nn.Sequential):
+    """Shallow ConvNet model from [2]_.
+
+    Parameters
+    ----------
+    in_chans : int
+        XXX
+
+    References
+    ----------
+    .. [2] 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=None,
+        n_filters_time=40,
+        filter_time_length=25,
+        n_filters_spat=40,
+        pool_time_length=75,
+        pool_time_stride=15,
+        final_conv_length=30,
+        conv_nonlin=square,
+        pool_mode="mean",
+        pool_nonlin=safe_log,
+        split_first_layer=True,
+        batch_norm=True,
+        batch_norm_alpha=0.1,
+        drop_prob=0.5,
+    ):
+        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.n_filters_time = n_filters_time
+        self.filter_time_length = filter_time_length
+        self.n_filters_spat = n_filters_spat
+        self.pool_time_length = pool_time_length
+        self.pool_time_stride = pool_time_stride
+        self.final_conv_length = final_conv_length
+        self.conv_nonlin = conv_nonlin
+        self.pool_mode = pool_mode
+        self.pool_nonlin = pool_nonlin
+        self.split_first_layer = split_first_layer
+        self.batch_norm = batch_norm
+        self.batch_norm_alpha = batch_norm_alpha
+        self.drop_prob = drop_prob
+
+        self.add_module("ensuredims", Ensure4d())
+        pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.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=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=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
+                ),
+            )
+        self.add_module("conv_nonlin_exp", Expression(self.conv_nonlin))
+        self.add_module(
+            "pool",
+            pool_class(
+                kernel_size=(self.pool_time_length, 1),
+                stride=(self.pool_time_stride, 1),
+            ),
+        )
+        self.add_module("pool_nonlin_exp", Expression(self.pool_nonlin))
+        self.add_module("drop", 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(
+                n_filters_conv,
+                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 paper...
+        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)
+        init.xavier_uniform_(self.conv_classifier.weight, gain=1)
+        init.constant_(self.conv_classifier.bias, 0)