braindecode 0.1.4.1__tar.gz → 0.1.4.2__tar.gz

{Braindecode-0.1.4.1 → Braindecode-0.1.4.2}/Braindecode.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.1
 Name: Braindecode
-Version: 0.1.4.1
+Version: 0.1.4.2
 Summary: A deep learning toolbox to decode raw time-domain EEG.
 Home-page: https://github.com/robintibor/braindecode
 Author: Robin Tibor Schirrmeister

{Braindecode-0.1.4.1 → Braindecode-0.1.4.2}/Braindecode.egg-info/SOURCES.txt

@@ -27,6 +27,7 @@ braindecode/experiments/stopcriteria.py
 braindecode/mne_ext/__init__.py
 braindecode/mne_ext/signalproc.py
 braindecode/models/__init__.py
+braindecode/models/deep4.py
 braindecode/models/shallow_fbcsp.py
 braindecode/torch_ext/__init__.py
 braindecode/torch_ext/constraints.py

{Braindecode-0.1.4.1 → Braindecode-0.1.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.1
 Name: Braindecode
-Version: 0.1.4.1
+Version: 0.1.4.2
 Summary: A deep learning toolbox to decode raw time-domain EEG.
 Home-page: https://github.com/robintibor/braindecode
 Author: Robin Tibor Schirrmeister

Braindecode-0.1.4.2/braindecode/models/deep4.py

@@ -0,0 +1,159 @@
+import numpy as np
+from torch import nn
+from torch.nn import init
+from torch.nn.functional import elu
+from braindecode.torch_ext.modules import Expression
+from braindecode.torch_ext.functions import identity
+from braindecode.torch_ext.util import np_to_var
+
+
+class Deep4Net(object):
+    def __init__(self, in_chans,
+                 n_classes,
+                 input_time_length,
+                 final_conv_length,
+                 n_filters_time=25,
+                 n_filters_spat=25,
+                 filter_time_length=10,
+                 pool_time_length=3,
+                 pool_time_stride=3,
+                 num_filters_2=50,
+                 filter_length_2=10,
+                 num_filters_3=100,
+                 filter_length_3=10,
+                 num_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):
+        if final_conv_length == 'auto':
+            assert input_time_length is not None
+
+        self.__dict__.update(locals())
+        del self.self
+
+    def create_network(self):
+        pool_class_dict = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)
+        first_pool_class = pool_class_dict[self.first_pool_mode]
+        later_pool_class = pool_class_dict[self.later_pool_mode]
+        model = nn.Sequential()
+        if self.split_first_layer:
+            model.add_module('dimshuffle',
+                             Expression(lambda x: x.permute(0, 3, 2, 1)))
+            model.add_module('conv_time', nn.Conv2d(1, self.n_filters_time,
+                                                    (
+                                                    self.filter_time_length, 1),
+                                                    stride=1, ))
+            model.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:
+            model.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:
+            model.add_module('bnorm',
+                             nn.BatchNorm2d(n_filters_conv,
+                                            momentum=self.batch_norm_alpha,
+                                            affine=True,
+                                            eps=1e-5),)
+        model.add_module('conv_nonlin', Expression(self.first_nonlin))
+        model.add_module('pool',
+                         first_pool_class(
+                             kernel_size=(self.pool_time_length, 1),
+                             stride=(self.pool_time_stride, 1)))
+        model.add_module('pool_nonlin', Expression(self.first_pool_nonlin))
+
+        def add_conv_pool_block(model, num_filters_before,
+                                num_filters, filter_length, block_nr):
+            suffix = '_{:d}'.format(block_nr)
+            model.add_module('drop' + suffix,
+                             nn.Dropout(p=self.drop_prob))
+            model.add_module('conv' + suffix.format(block_nr),
+                             nn.Conv2d(num_filters_before, num_filters,
+                                       (filter_length, 1),
+                                       stride=1, bias=not self.batch_norm))
+            if self.batch_norm:
+                model.add_module('bnorm' + suffix,
+                             nn.BatchNorm2d(num_filters,
+                                            momentum=self.batch_norm_alpha,
+                                            affine=True,
+                                            eps=1e-5))
+            model.add_module('nonlin' + suffix,
+                             Expression(self.later_nonlin))
+
+            model.add_module('pool' + suffix,
+                             later_pool_class(
+                                 kernel_size=(self.pool_time_length, 1),
+                                 stride=(self.pool_time_stride, 1)))
+            model.add_module('pool_nonlin' + suffix,
+                             Expression(self.later_pool_nonlin))
+
+        add_conv_pool_block(model, n_filters_conv, self.num_filters_2,
+                            self.filter_length_2, 2)
+        add_conv_pool_block(model, self.num_filters_2, self.num_filters_3,
+                            self.filter_length_3, 3)
+        add_conv_pool_block(model, self.num_filters_3, self.num_filters_4,
+                            self.filter_length_4, 4)
+
+
+        if self.final_conv_length == 'auto':
+            out = model(np_to_var(np.ones(
+                (1, self.in_chans, self.input_time_length,1),
+                dtype=np.float32)))
+            n_out_time = out.cpu().data.numpy().shape[2]
+            self.final_conv_length = n_out_time
+        model.add_module('conv_classifier',
+                             nn.Conv2d(self.num_filters_4, self.n_classes,
+                                       (self.final_conv_length, 1), bias=True))
+        model.add_module('softmax', nn.LogSoftmax())
+        # remove empty dim at end and potentially remove empty time dim
+        # do not just use squeeze as we never want to remove first dim
+        def squeeze_output(x):
+            assert x.size()[3] == 1
+            x = x[:,:,:,0]
+            if x.size()[2] == 1:
+                x = x[:,:,0]
+            return x
+        model.add_module('squeeze',  Expression(squeeze_output))
+
+        # Initialization, xavier is same as in our paper...
+        # was default from lasagne
+        init.xavier_uniform(model.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(model.conv_time.bias, 0)
+        if self.split_first_layer:
+            init.xavier_uniform(model.conv_spat.weight, gain=1)
+            if not self.batch_norm:
+                init.constant(model.conv_spat.bias, 0)
+        param_dict = dict(list(model.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(model.conv_classifier.weight, gain=1)
+        init.constant(model.conv_classifier.bias, 0)
+
+        return model

{Braindecode-0.1.4.1 → Braindecode-0.1.4.2}/braindecode/models/shallow_fbcsp.py

@@ -8,7 +8,6 @@ from braindecode.torch_ext.util import np_to_var
 
 
 class ShallowFBCSPNet(object):
-    # TODO: auto final dense length for shallow
     def __init__(self, in_chans,
                  n_classes,
                  input_time_length=None,
@@ -25,7 +24,7 @@ class ShallowFBCSPNet(object):
                  batch_norm=True,
                  batch_norm_alpha=0.1,
                  drop_prob=0.5):
-        if final_conv_length == 'full':
+        if final_conv_length == 'auto':
             assert input_time_length is not None
         self.__dict__.update(locals())
         del self.self
@@ -92,6 +91,9 @@ class ShallowFBCSPNet(object):
             init.xavier_uniform(model.conv_spat.weight, gain=1)
             if not self.batch_norm:
                 init.constant(model.conv_spat.bias, 0)
+            else:
+                init.constant(model.bnorm.weight, 1)
+                init.constant(model.bnorm.bias, 0)
         init.xavier_uniform(model.conv_classifier.weight, gain=1)
         init.constant(model.conv_classifier.bias, 0)
 

{Braindecode-0.1.4.1 → Braindecode-0.1.4.2}/braindecode/torch_ext/functions.py

@@ -7,3 +7,7 @@ def square(x):
 def safe_log(x, eps=1e-6):
     """ Prevents log(0) by using max of eps and given x."""
     return th.log(th.clamp(x, min=eps))
+
+
+def identity(x):
+    return x