[WIP] Spatial eegnet

.gitignore

@@ -158,4 +158,8 @@ dmypy.json
 **/log/
 
 # Mac OS
-.DS_Store
+.DS_Store
+
+results
+eeg_data
+trial_*.conf

benchmarks/MOABB/extra-requirements.txt

@@ -1,6 +1,7 @@
-mne==1.6.1
+mne
 moabb
 orion
 orion[profet]
 scikit-learn
+torch_geometric
 torchinfo

benchmarks/MOABB/hparams/MotorImagery/BNCI2014001/SpatialEEGNet.yaml

@@ -0,0 +1,187 @@
+seed: 1234
+__set_torchseed: !apply:torch.manual_seed [!ref <seed>]
+
+# DIRECTORIES
+data_folder:
+  !PLACEHOLDER #'/path/to/dataset'. The dataset will be automatically downloaded in this folder
+
+
+cached_data_folder: !PLACEHOLDER #'path/to/pickled/dataset'
+
+
+output_folder: !PLACEHOLDER
+  #'path/to/results'
+
+  # DATASET HPARS
+  # Defining the MOABB dataset.
+
+
+datasets:
+  - !new:moabb.datasets.BNCI2014001
+save_prepared_dataset: True # set to True if you want to save the prepared dataset as a pkl file to load and use afterwards
+data_iterator_name: !PLACEHOLDER
+target_subject_idx: !PLACEHOLDER
+target_session_idx: !PLACEHOLDER
+events_to_load: null # all events will be loaded
+original_sample_rate: 250 # Original sampling rate provided by dataset authors
+sample_rate: 125 # Target sampling rate (Hz)
+# band-pass filtering cut-off frequencies
+fmin: 0.13 # orion_step1: --fmin~"uniform(0.1, 5, precision=2)"
+fmax: 46.0 # @orion_step1: --fmax~"uniform(20.0, 50.0, precision=3)"
+n_classes: 4
+# tmin, tmax respect to stimulus onset that define the interval attribute of the dataset class
+# trial begins (0 s), cue (2 s, 1.25 s long); each trial is 6 s long
+# dataset interval starts from 2
+# -->tmin tmax are referred to this start value (e.g., tmin=0.5 corresponds to 2.5 s)
+tmin: 0.
+tmax: 4.0 # @orion_step1: --tmax~"uniform(1.0, 4.0, precision=2)"
+n_steps_channel_selection: null
+T: !apply:math.ceil
+  - !ref <sample_rate> * (<tmax> - <tmin>)
+C: 22
+# We here specify how to perfom test:
+# - If test_with: 'last' we perform test with the latest model.
+# - if test_with: 'best, we perform test with the best model (according to the metric specified in test_key)
+# The variable avg_models can be used to average the parameters of the last (or best) N saved models before testing.
+# This can have a regularization effect. If avg_models: 1, the last (or best) model is used directly.
+test_with: "last" # 'last' or 'best'
+test_key: "acc" # Possible opts: "loss", "f1", "auc", "acc"
+
+# METRICS
+f1: !name:sklearn.metrics.f1_score
+  average: "macro"
+acc: !name:sklearn.metrics.balanced_accuracy_score
+cm: !name:sklearn.metrics.confusion_matrix
+metrics:
+  f1: !ref <f1>
+  acc: !ref <acc>
+  cm: !ref <cm>
+# TRAINING HPARS
+n_train_examples: 100 # it will be replaced in the train script
+# checkpoints to average
+avg_models: 10 # @orion_step1: --avg_models~"uniform(1, 15,discrete=True)"
+number_of_epochs: 862 # @orion_step1: --number_of_epochs~"uniform(500, 1000, discrete=True)"
+lr: 0.0001 # orion_step1: --lr~"choices([0.01, 0.005, 0.001, 0.0005, 0.0001])"
+# Learning rate scheduling (cyclic learning rate is used here)
+max_lr: !ref <lr> # Upper bound of the cycle (max value of the lr)
+base_lr: 0.00000001 # Lower bound in the cycle (min value of the lr)
+step_size_multiplier: 5 #from 2 to 8
+step_size: !apply:round
+  - !ref <step_size_multiplier> * <n_train_examples> / <batch_size>
+lr_annealing: !new:speechbrain.nnet.schedulers.CyclicLRScheduler
+  base_lr: !ref <base_lr>
+  max_lr: !ref <max_lr>
+  step_size: !ref <step_size>
+label_smoothing: 0.0
+loss: !name:speechbrain.nnet.losses.nll_loss
+  label_smoothing: !ref <label_smoothing>
+optimizer: !name:torch.optim.Adam
+  lr: !ref <lr>
+epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter # epoch counter
+  limit: !ref <number_of_epochs>
+batch_size_exponent: 4 # @orion_step1: --batch_size_exponent~"uniform(4, 6,discrete=True)"
+batch_size: !ref 2 ** <batch_size_exponent>
+valid_ratio: 0.2
+
+# DATA AUGMENTATION
+# cutcat (disabled when min_num_segments=max_num_segments=1)
+max_num_segments: 3 # orion_step2: --max_num_segments~"uniform(2, 6, discrete=True)"
+cutcat: !new:speechbrain.augment.time_domain.CutCat
+  min_num_segments: 2
+  max_num_segments: !ref <max_num_segments>
+# random amplitude gain between 0.5-1.5 uV (disabled when amp_delta=0.)
+amp_delta: 0.01742 # orion_step2: --amp_delta~"uniform(0.0, 0.5)"
+rand_amp: !new:speechbrain.augment.time_domain.RandAmp
+  amp_low: !ref 1 - <amp_delta>
+  amp_high: !ref 1 + <amp_delta>
+# random shifts between -300 ms to 300 ms (disabled when shift_delta=0.)
+shift_delta_: 1 # orion_step2: --shift_delta_~"uniform(0, 25, discrete=True)"
+shift_delta: !ref 1e-2 * <shift_delta_> # 0.250 # 0.-0.25 with steps of 0.01
+min_shift: !apply:math.floor
+  - !ref 0 - <sample_rate> * <shift_delta>
+max_shift: !apply:math.floor
+  - !ref 0 + <sample_rate> * <shift_delta>
+time_shift: !new:speechbrain.augment.freq_domain.RandomShift
+  min_shift: !ref <min_shift>
+  max_shift: !ref <max_shift>
+  dim: 1
+# injection of gaussian white noise
+snr_white_low: 15.0 # orion_step2: --snr_white_low~"uniform(0.0, 15, precision=2)"
+snr_white_delta: 19.1 # orion_step2: --snr_white_delta~"uniform(5.0, 20.0, precision=3)"
+snr_white_high: !ref <snr_white_low> + <snr_white_delta>
+add_noise_white: !new:speechbrain.augment.time_domain.AddNoise
+  snr_low: !ref <snr_white_low>
+  snr_high: !ref <snr_white_high>
+position_noise_sigma: 0.01 # @orion_step1: --position_noise_sigma~"uniform(0.0, 0.1, precision=4)"
+
+repeat_augment: 1 # @orion_step1: --repeat_augment 0
+graph_augment: !new:torch.nn.Sequential
+  - !new:utils.graph_iterators.PositionNoise
+    sigma: !ref <position_noise_sigma>
+  - !new:utils.graph_iterators.NodeDrop
+    choose_k: 17
+augment: !new:speechbrain.augment.augmenter.Augmenter
+  parallel_augment: True
+  concat_original: True
+  parallel_augment_fixed_bs: True
+  repeat_augment: !ref <repeat_augment>
+  shuffle_augmentations: True
+  min_augmentations: 4
+  max_augmentations: 4
+  augmentations:
+    [!ref <cutcat>, !ref <rand_amp>, !ref <time_shift>, !ref <add_noise_white>]
+
+# DATA NORMALIZATION
+dims_to_normalize: 1 # 1 (time) or 2 (EEG channels)
+normalize: !name:speechbrain.processing.signal_processing.mean_std_norm
+  dims: !ref <dims_to_normalize>
+pos_normalize: True
+
+# MODEL
+projection_dim: 22 # @orion_step1: --projection_dim~"uniform(4, 44, discrete=True)"
+spatial_focus_tau: 0.3333 # @orion_step1: --spatial_focus_tau~"uniform(0.01, 1.0, precision=4)"
+
+input_shape: [null, !ref <T>, !ref <projection_dim>, null]
+cnn_temporal_kernels: 61 # @orion_step1: --cnn_temporal_kernels~"uniform(32, 72,discrete=True)"
+cnn_temporal_kernelsize: 51 # @orion_step1: --cnn_temporal_kernelsize~"uniform(24, 62,discrete=True)"
+# depth multiplier for the spatial depthwise conv. layer
+cnn_spatial_depth_multiplier: 4 # @orion_step1: --cnn_spatial_depth_multiplier~"uniform(1, 4,discrete=True)"
+cnn_spatial_max_norm: 1. # kernel max-norm constaint of the spatial depthwise conv. layer
+cnn_spatial_pool: 4
+cnn_septemporal_depth_multiplier: 1 # depth multiplier for the separable temporal conv. layer
+cnn_septemporal_point_kernels_ratio_: 7 # @orion_step1: --cnn_septemporal_point_kernels_ratio_~"uniform(0, 8, discrete=True)"
+cnn_septemporal_point_kernels_ratio: !ref <cnn_septemporal_point_kernels_ratio_> / 4
+## number of temporal filters in the separable temporal conv. layer
+cnn_septemporal_point_kernels_: !ref <cnn_temporal_kernels> * <cnn_spatial_depth_multiplier> * <cnn_septemporal_depth_multiplier>
+cnn_septemporal_point_kernels: !apply:math.ceil
+  - !ref <cnn_septemporal_point_kernels_ratio> * <cnn_septemporal_point_kernels_> + 1
+cnn_septemporal_kernelsize_: 15 # @orion_step1: --cnn_septemporal_kernelsize_~"uniform(3, 24,discrete=True)"
+max_cnn_spatial_pool: 4
+cnn_septemporal_kernelsize: !apply:round
+  - !ref <cnn_septemporal_kernelsize_> * <max_cnn_spatial_pool> / <cnn_spatial_pool>
+cnn_septemporal_pool: 7 # @orion_step1: --cnn_septemporal_pool~"uniform(1, 8,discrete=True)"
+cnn_pool_type: "avg"
+dense_max_norm: 0.25 # kernel max-norm constaint of the dense layer
+dropout: 0.008464 # @orion_step1: --dropout~"uniform(0.0, 0.5)"
+activation_type: "elu"
+
+model: !new:models.SpatialEEGNet.SpatialEEGNet
+  input_shape: !ref <input_shape>
+  cnn_temporal_kernels: !ref <cnn_temporal_kernels>
+  cnn_temporal_kernelsize: [!ref <cnn_temporal_kernelsize>, 1]
+  cnn_spatial_depth_multiplier: !ref <cnn_spatial_depth_multiplier>
+  cnn_spatial_max_norm: !ref <cnn_spatial_max_norm>
+  cnn_spatial_pool: [!ref <cnn_spatial_pool>, 1]
+  cnn_septemporal_depth_multiplier: !ref <cnn_septemporal_depth_multiplier>
+  cnn_septemporal_point_kernels: !ref <cnn_septemporal_point_kernels>
+  cnn_septemporal_kernelsize: [!ref <cnn_septemporal_kernelsize>, 1]
+  cnn_septemporal_pool: [!ref <cnn_septemporal_pool>, 1]
+  cnn_pool_type: !ref <cnn_pool_type>
+  activation_type: !ref <activation_type>
+  spatial_focus: !new:models.SpatialEEGNet.SpatialFocus
+    projection_dim: !ref <projection_dim>
+    position_dim: 3
+    tau: !ref <spatial_focus_tau>
+  dense_max_norm: !ref <dense_max_norm>
+  dropout: !ref <dropout>
+  dense_n_neurons: !ref <n_classes>

benchmarks/MOABB/hparams/MotorImagery/BNCI2014001/SpatialEEGNet_cross.yaml

@@ -0,0 +1,187 @@
+seed: 1234
+__set_torchseed: !apply:torch.manual_seed [!ref <seed>]
+
+# DIRECTORIES
+data_folder:
+  !PLACEHOLDER #'/path/to/dataset'. The dataset will be automatically downloaded in this folder
+
+
+cached_data_folder: !PLACEHOLDER #'path/to/pickled/dataset'
+
+
+output_folder: !PLACEHOLDER
+  #'path/to/results'
+
+  # DATASET HPARS
+  # Defining the MOABB dataset.
+
+
+datasets:
+  - !new:moabb.datasets.BNCI2014001
+save_prepared_dataset: True # set to True if you want to save the prepared dataset as a pkl file to load and use afterwards
+data_iterator_name: !PLACEHOLDER
+target_subject_idx: !PLACEHOLDER
+target_session_idx: !PLACEHOLDER
+events_to_load: null # all events will be loaded
+original_sample_rate: 250 # Original sampling rate provided by dataset authors
+sample_rate: 125 # Target sampling rate (Hz)
+# band-pass filtering cut-off frequencies
+fmin: 0.13 # orion_step1: --fmin~"uniform(0.1, 5, precision=2)"
+fmax: 46.0 # @orion_step1: --fmax~"uniform(20.0, 50.0, precision=3)"
+n_classes: 4
+# tmin, tmax respect to stimulus onset that define the interval attribute of the dataset class
+# trial begins (0 s), cue (2 s, 1.25 s long); each trial is 6 s long
+# dataset interval starts from 2
+# -->tmin tmax are referred to this start value (e.g., tmin=0.5 corresponds to 2.5 s)
+tmin: 0.
+tmax: 4.0 # @orion_step1: --tmax~"uniform(1.0, 4.0, precision=2)"
+n_steps_channel_selection: null
+T: !apply:math.ceil
+  - !ref <sample_rate> * (<tmax> - <tmin>)
+C: 22
+# We here specify how to perfom test:
+# - If test_with: 'last' we perform test with the latest model.
+# - if test_with: 'best, we perform test with the best model (according to the metric specified in test_key)
+# The variable avg_models can be used to average the parameters of the last (or best) N saved models before testing.
+# This can have a regularization effect. If avg_models: 1, the last (or best) model is used directly.
+test_with: "last" # 'last' or 'best'
+test_key: "acc" # Possible opts: "loss", "f1", "auc", "acc"
+
+# METRICS
+f1: !name:sklearn.metrics.f1_score
+  average: "macro"
+acc: !name:sklearn.metrics.balanced_accuracy_score
+cm: !name:sklearn.metrics.confusion_matrix
+metrics:
+  f1: !ref <f1>
+  acc: !ref <acc>
+  cm: !ref <cm>
+# TRAINING HPARS
+n_train_examples: 100 # it will be replaced in the train script
+# checkpoints to average
+avg_models: 10 # @orion_step1: --avg_models~"uniform(1, 15,discrete=True)"
+number_of_epochs: 100 # orion_step1: --number_of_epochs~"uniform(500, 1000, discrete=True)"
+lr: 0.0001 # orion_step1: --lr~"choices([0.01, 0.005, 0.001, 0.0005, 0.0001])"
+# Learning rate scheduling (cyclic learning rate is used here)
+max_lr: !ref <lr> # Upper bound of the cycle (max value of the lr)
+base_lr: 0.00000001 # Lower bound in the cycle (min value of the lr)
+step_size_multiplier: 5 #from 2 to 8
+step_size: !apply:round
+  - !ref <step_size_multiplier> * <n_train_examples> / <batch_size>
+lr_annealing: !new:speechbrain.nnet.schedulers.CyclicLRScheduler
+  base_lr: !ref <base_lr>
+  max_lr: !ref <max_lr>
+  step_size: !ref <step_size>
+label_smoothing: 0.0
+loss: !name:speechbrain.nnet.losses.nll_loss
+  label_smoothing: !ref <label_smoothing>
+optimizer: !name:torch.optim.Adam
+  lr: !ref <lr>
+epoch_counter: !new:speechbrain.utils.epoch_loop.EpochCounter # epoch counter
+  limit: !ref <number_of_epochs>
+batch_size_exponent: 4 # @orion_step1: --batch_size_exponent~"uniform(4, 6,discrete=True)"
+batch_size: !ref 2 ** <batch_size_exponent>
+valid_ratio: 0.2
+
+# DATA AUGMENTATION
+# cutcat (disabled when min_num_segments=max_num_segments=1)
+max_num_segments: 3 # @orion_step2: --max_num_segments~"uniform(2, 6, discrete=True)"
+cutcat: !new:speechbrain.augment.time_domain.CutCat
+  min_num_segments: 2
+  max_num_segments: !ref <max_num_segments>
+# random amplitude gain between 0.5-1.5 uV (disabled when amp_delta=0.)
+amp_delta: 0.01742 # @orion_step2: --amp_delta~"uniform(0.0, 0.5)"
+rand_amp: !new:speechbrain.augment.time_domain.RandAmp
+  amp_low: !ref 1 - <amp_delta>
+  amp_high: !ref 1 + <amp_delta>
+# random shifts between -300 ms to 300 ms (disabled when shift_delta=0.)
+shift_delta_: 1 # @orion_step2: --shift_delta_~"uniform(0, 25, discrete=True)"
+shift_delta: !ref 1e-2 * <shift_delta_> # 0.250 # 0.-0.25 with steps of 0.01
+min_shift: !apply:math.floor
+  - !ref 0 - <sample_rate> * <shift_delta>
+max_shift: !apply:math.floor
+  - !ref 0 + <sample_rate> * <shift_delta>
+time_shift: !new:speechbrain.augment.freq_domain.RandomShift
+  min_shift: !ref <min_shift>
+  max_shift: !ref <max_shift>
+  dim: 1
+# injection of gaussian white noise
+snr_white_low: 15.0 # @orion_step2: --snr_white_low~"uniform(0.0, 15, precision=2)"
+snr_white_delta: 19.1 # @orion_step2: --snr_white_delta~"uniform(5.0, 20.0, precision=3)"
+snr_white_high: !ref <snr_white_low> + <snr_white_delta>
+add_noise_white: !new:speechbrain.augment.time_domain.AddNoise
+  snr_low: !ref <snr_white_low>
+  snr_high: !ref <snr_white_high>
+position_noise_sigma: 0.01 # @orion_step1: --position_noise_sigma~"uniform(0.0, 0.1, precision=4)"
+
+repeat_augment: 1 # @orion_step1: --repeat_augment 0
+graph_augment: !new:torch.nn.Sequential
+  - !new:utils.graph_iterators.PositionNoise
+    sigma: !ref <position_noise_sigma>
+  # - !new:utils.graph_iterators.NodeDrop
+  #   choose_k: 17
+augment: !new:speechbrain.augment.augmenter.Augmenter
+  parallel_augment: True
+  concat_original: True
+  parallel_augment_fixed_bs: True
+  repeat_augment: !ref <repeat_augment>
+  shuffle_augmentations: True
+  min_augmentations: 4
+  max_augmentations: 4
+  augmentations:
+    [!ref <cutcat>, !ref <rand_amp>, !ref <time_shift>, !ref <add_noise_white>]
+
+# DATA NORMALIZATION
+dims_to_normalize: 1 # 1 (time) or 2 (EEG channels)
+normalize: !name:speechbrain.processing.signal_processing.mean_std_norm
+  dims: !ref <dims_to_normalize>
+pos_normalize: True
+
+# MODEL
+projection_dim: 22 # @orion_step1: --projection_dim~"uniform(4, 44, discrete=True)"
+spatial_focus_tau: 0.3333 # @orion_step1: --spatial_focus_tau~"uniform(0.001, 1.0, precision=4)"
+
+input_shape: [null, !ref <T>, !ref <projection_dim>, null]
+cnn_temporal_kernels: 61 # @orion_step1: --cnn_temporal_kernels~"uniform(32, 72,discrete=True)"
+cnn_temporal_kernelsize: 51 # @orion_step1: --cnn_temporal_kernelsize~"uniform(24, 62,discrete=True)"
+# depth multiplier for the spatial depthwise conv. layer
+cnn_spatial_depth_multiplier: 4 # @orion_step1: --cnn_spatial_depth_multiplier~"uniform(1, 4,discrete=True)"
+cnn_spatial_max_norm: 1. # kernel max-norm constaint of the spatial depthwise conv. layer
+cnn_spatial_pool: 4
+cnn_septemporal_depth_multiplier: 1 # depth multiplier for the separable temporal conv. layer
+cnn_septemporal_point_kernels_ratio_: 7 # @orion_step1: --cnn_septemporal_point_kernels_ratio_~"uniform(0, 8, discrete=True)"
+cnn_septemporal_point_kernels_ratio: !ref <cnn_septemporal_point_kernels_ratio_> / 4
+## number of temporal filters in the separable temporal conv. layer
+cnn_septemporal_point_kernels_: !ref <cnn_temporal_kernels> * <cnn_spatial_depth_multiplier> * <cnn_septemporal_depth_multiplier>
+cnn_septemporal_point_kernels: !apply:math.ceil
+  - !ref <cnn_septemporal_point_kernels_ratio> * <cnn_septemporal_point_kernels_> + 1
+cnn_septemporal_kernelsize_: 15 # @orion_step1: --cnn_septemporal_kernelsize_~"uniform(3, 24,discrete=True)"
+max_cnn_spatial_pool: 4
+cnn_septemporal_kernelsize: !apply:round
+  - !ref <cnn_septemporal_kernelsize_> * <max_cnn_spatial_pool> / <cnn_spatial_pool>
+cnn_septemporal_pool: 7 # @orion_step1: --cnn_septemporal_pool~"uniform(1, 8,discrete=True)"
+cnn_pool_type: "avg"
+dense_max_norm: 0.25 # kernel max-norm constaint of the dense layer
+dropout: 0.008464 # @orion_step1: --dropout~"uniform(0.0, 0.5)"
+activation_type: "elu"
+
+model: !new:models.SpatialEEGNet.SpatialEEGNet
+  input_shape: !ref <input_shape>
+  cnn_temporal_kernels: !ref <cnn_temporal_kernels>
+  cnn_temporal_kernelsize: [!ref <cnn_temporal_kernelsize>, 1]
+  cnn_spatial_depth_multiplier: !ref <cnn_spatial_depth_multiplier>
+  cnn_spatial_max_norm: !ref <cnn_spatial_max_norm>
+  cnn_spatial_pool: [!ref <cnn_spatial_pool>, 1]
+  cnn_septemporal_depth_multiplier: !ref <cnn_septemporal_depth_multiplier>
+  cnn_septemporal_point_kernels: !ref <cnn_septemporal_point_kernels>
+  cnn_septemporal_kernelsize: [!ref <cnn_septemporal_kernelsize>, 1]
+  cnn_septemporal_pool: [!ref <cnn_septemporal_pool>, 1]
+  cnn_pool_type: !ref <cnn_pool_type>
+  activation_type: !ref <activation_type>
+  spatial_focus: !new:models.SpatialEEGNet.SpatialFocus
+    projection_dim: !ref <projection_dim>
+    position_dim: 3
+    tau: !ref <spatial_focus_tau>
+  dense_max_norm: !ref <dense_max_norm>
+  dropout: !ref <dropout>
+  dense_n_neurons: !ref <n_classes>