def run_exp( data_folders, n_recordings, sensor_types, n_chans, max_recording_mins, sec_to_cut, duration_recording_mins, test_recording_mins, max_abs_val, sampling_freq, divisor, test_on_eval, n_folds, i_test_fold, shuffle, model_name, n_start_chans, n_chan_factor, input_time_length, final_conv_length, model_constraint, init_lr, batch_size, max_epochs, cuda, ): import torch.backends.cudnn as cudnn cudnn.benchmark = True preproc_functions = [] preproc_functions.append(lambda data, fs: ( data[:, int(sec_to_cut * fs):-int(sec_to_cut * fs)], fs)) preproc_functions.append(lambda data, fs: (data[:, :int( duration_recording_mins * 60 * fs)], fs)) if max_abs_val is not None: preproc_functions.append( lambda data, fs: (np.clip(data, -max_abs_val, max_abs_val), fs)) preproc_functions.append(lambda data, fs: (resampy.resample( data, fs, sampling_freq, axis=1, filter='kaiser_fast'), sampling_freq)) if divisor is not None: preproc_functions.append(lambda data, fs: (data / divisor, fs)) dataset = DiagnosisSet(n_recordings=n_recordings, max_recording_mins=max_recording_mins, preproc_functions=preproc_functions, data_folders=data_folders, train_or_eval='train', sensor_types=sensor_types) if test_on_eval: if test_recording_mins is None: test_recording_mins = duration_recording_mins test_preproc_functions = copy(preproc_functions) test_preproc_functions[1] = lambda data, fs: (data[:, :int( test_recording_mins * 60 * fs)], fs) test_dataset = DiagnosisSet(n_recordings=n_recordings, max_recording_mins=None, preproc_functions=test_preproc_functions, data_folders=data_folders, train_or_eval='eval', sensor_types=sensor_types) X, y = dataset.load() max_shape = np.max([list(x.shape) for x in X], axis=0) assert max_shape[1] == int(duration_recording_mins * sampling_freq * 60) if test_on_eval: test_X, test_y = test_dataset.load() max_shape = np.max([list(x.shape) for x in test_X], axis=0) assert max_shape[1] == int(test_recording_mins * sampling_freq * 60) if not test_on_eval: splitter = TrainValidTestSplitter(n_folds, i_test_fold, shuffle=shuffle) train_set, valid_set, test_set = splitter.split(X, y) else: splitter = TrainValidSplitter(n_folds, i_valid_fold=i_test_fold, shuffle=shuffle) train_set, valid_set = splitter.split(X, y) test_set = SignalAndTarget(test_X, test_y) del test_X, test_y del X, y # shouldn't be necessary, but just to make sure set_random_seeds(seed=20170629, cuda=cuda) n_classes = 2 if model_name == 'shallow': model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, final_conv_length=final_conv_length).create_network() elif model_name == 'deep': model = Deep4Net(n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, stride_before_pool=True).create_network() elif (model_name == 'deep_smac'): if model_name == 'deep_smac': do_batch_norm = False else: assert model_name == 'deep_smac_bnorm' do_batch_norm = True double_time_convs = False drop_prob = 0.244445 filter_length_2 = 12 filter_length_3 = 14 filter_length_4 = 12 filter_time_length = 21 final_conv_length = 1 first_nonlin = elu first_pool_mode = 'mean' first_pool_nonlin = identity later_nonlin = elu later_pool_mode = 'mean' later_pool_nonlin = identity n_filters_factor = 1.679066 n_filters_start = 32 pool_time_length = 1 pool_time_stride = 2 split_first_layer = True n_chan_factor = n_filters_factor n_start_chans = n_filters_start model = Deep4Net(n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, batch_norm=do_batch_norm, double_time_convs=double_time_convs, drop_prob=drop_prob, filter_length_2=filter_length_2, filter_length_3=filter_length_3, filter_length_4=filter_length_4, filter_time_length=filter_time_length, first_nonlin=first_nonlin, first_pool_mode=first_pool_mode, first_pool_nonlin=first_pool_nonlin, later_nonlin=later_nonlin, later_pool_mode=later_pool_mode, later_pool_nonlin=later_pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, stride_before_pool=True).create_network() elif model_name == 'shallow_smac': conv_nonlin = identity do_batch_norm = True drop_prob = 0.328794 filter_time_length = 56 final_conv_length = 22 n_filters_spat = 73 n_filters_time = 24 pool_mode = 'max' pool_nonlin = identity pool_time_length = 84 pool_time_stride = 3 split_first_layer = True model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_filters_time, n_filters_spat=n_filters_spat, input_time_length=input_time_length, final_conv_length=final_conv_length, conv_nonlin=conv_nonlin, batch_norm=do_batch_norm, drop_prob=drop_prob, filter_time_length=filter_time_length, pool_mode=pool_mode, pool_nonlin=pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, ).create_network() elif model_name == 'linear': model = nn.Sequential() model.add_module("conv_classifier", nn.Conv2d(n_chans, n_classes, (600, 1))) model.add_module('softmax', nn.LogSoftmax()) model.add_module('squeeze', Expression(lambda x: x.squeeze(3))) else: assert False, "unknown model name {:s}".format(model_name) to_dense_prediction_model(model) log.info("Model:\n{:s}".format(str(model))) if cuda: model.cuda() # determine output size test_input = np_to_var( np.ones((2, n_chans, input_time_length, 1), dtype=np.float32)) if cuda: test_input = test_input.cuda() log.info("In shape: {:s}".format(str(test_input.cpu().data.numpy().shape))) out = model(test_input) log.info("Out shape: {:s}".format(str(out.cpu().data.numpy().shape))) n_preds_per_input = out.cpu().data.numpy().shape[2] log.info("{:d} predictions per input/trial".format(n_preds_per_input)) iterator = CropsFromTrialsIterator(batch_size=batch_size, input_time_length=input_time_length, n_preds_per_input=n_preds_per_input) optimizer = optim.Adam(model.parameters(), lr=init_lr) loss_function = lambda preds, targets: F.nll_loss( th.mean(preds, dim=2, keepdim=False), targets) if model_constraint is not None: assert model_constraint == 'defaultnorm' model_constraint = MaxNormDefaultConstraint() monitors = [ LossMonitor(), MisclassMonitor(col_suffix='sample_misclass'), CroppedDiagnosisMonitor(input_time_length, n_preds_per_input), RuntimeMonitor(), ] stop_criterion = MaxEpochs(max_epochs) batch_modifier = None run_after_early_stop = True exp = Experiment(model, train_set, valid_set, test_set, iterator, loss_function, optimizer, model_constraint, monitors, stop_criterion, remember_best_column='valid_misclass', run_after_early_stop=run_after_early_stop, batch_modifier=batch_modifier, cuda=cuda) exp.run() return exp
def run_exp(test_on_eval, sensor_types, n_chans, max_recording_mins, test_recording_mins, n_recordings, sec_to_cut_at_start, sec_to_cut_at_end, duration_recording_mins, max_abs_val, clip_before_resample, sampling_freq, divisor, n_folds, i_test_fold, shuffle, merge_train_valid, model_name, n_start_chans, n_chan_factor, input_time_length, final_conv_length, stride_before_pool, optimizer, learning_rate, weight_decay, scheduler, model_constraint, batch_size, max_epochs, log_dir, only_return_exp, np_th_seed): cuda = True if ('smac' in model_name) and (input_time_length is None): input_time_length = 12000 fix_input_length_for_smac = True else: fix_input_length_for_smac = False set_random_seeds(seed=np_th_seed, cuda=cuda) n_classes = 2 if model_name == 'shallow': model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, final_conv_length=final_conv_length).create_network() elif model_name == 'deep': model = Deep4Net( n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, stride_before_pool=stride_before_pool).create_network() elif (model_name == 'deep_smac') or (model_name == 'deep_smac_bnorm'): if model_name == 'deep_smac': do_batch_norm = False else: assert model_name == 'deep_smac_bnorm' do_batch_norm = True double_time_convs = False drop_prob = 0.244445 filter_length_2 = 12 filter_length_3 = 14 filter_length_4 = 12 filter_time_length = 21 final_conv_length = 1 first_nonlin = elu first_pool_mode = 'mean' first_pool_nonlin = identity later_nonlin = elu later_pool_mode = 'mean' later_pool_nonlin = identity n_filters_factor = 1.679066 n_filters_start = 32 pool_time_length = 1 pool_time_stride = 2 split_first_layer = True n_chan_factor = n_filters_factor n_start_chans = n_filters_start model = Deep4Net(n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, batch_norm=do_batch_norm, double_time_convs=double_time_convs, drop_prob=drop_prob, filter_length_2=filter_length_2, filter_length_3=filter_length_3, filter_length_4=filter_length_4, filter_time_length=filter_time_length, first_nonlin=first_nonlin, first_pool_mode=first_pool_mode, first_pool_nonlin=first_pool_nonlin, later_nonlin=later_nonlin, later_pool_mode=later_pool_mode, later_pool_nonlin=later_pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, stride_before_pool=True).create_network() elif model_name == 'shallow_smac': conv_nonlin = identity do_batch_norm = True drop_prob = 0.328794 filter_time_length = 56 final_conv_length = 22 n_filters_spat = 73 n_filters_time = 24 pool_mode = 'max' pool_nonlin = identity pool_time_length = 84 pool_time_stride = 3 split_first_layer = True model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_filters_time, n_filters_spat=n_filters_spat, input_time_length=input_time_length, final_conv_length=final_conv_length, conv_nonlin=conv_nonlin, batch_norm=do_batch_norm, drop_prob=drop_prob, filter_time_length=filter_time_length, pool_mode=pool_mode, pool_nonlin=pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, ).create_network() elif model_name == 'deep_smac_new': from torch.nn.functional import elu, relu, relu6, tanh from braindecode.torch_ext.functions import identity, square, safe_log n_filters_factor = 1.9532637176784269 n_filters_start = 61 deep_kwargs = { "batch_norm": False, "double_time_convs": False, "drop_prob": 0.3622676569047184, "filter_length_2": 9, "filter_length_3": 6, "filter_length_4": 10, "filter_time_length": 17, "final_conv_length": 5, "first_nonlin": elu, "first_pool_mode": "max", "first_pool_nonlin": identity, "later_nonlin": relu6, "later_pool_mode": "max", "later_pool_nonlin": identity, "n_filters_time": n_filters_start, "n_filters_spat": n_filters_start, "n_filters_2": int(n_filters_start * n_filters_factor), "n_filters_3": int(n_filters_start * (n_filters_factor**2.0)), "n_filters_4": int(n_filters_start * (n_filters_factor**3.0)), "pool_time_length": 1, "pool_time_stride": 4, "split_first_layer": True, "stride_before_pool": True, } model = Deep4Net(n_chans, n_classes, input_time_length=input_time_length, **deep_kwargs).create_network() elif model_name == 'shallow_smac_new': from torch.nn.functional import elu, relu, relu6, tanh from braindecode.torch_ext.functions import identity, square, safe_log shallow_kwargs = { "conv_nonlin": square, "batch_norm": True, "drop_prob": 0.10198630723385381, "filter_time_length": 51, "final_conv_length": 1, "n_filters_spat": 200, "n_filters_time": 76, "pool_mode": "max", "pool_nonlin": safe_log, "pool_time_length": 139, "pool_time_stride": 49, "split_first_layer": True, } model = ShallowFBCSPNet(in_chans=n_chans, n_classes=n_classes, input_time_length=input_time_length, **shallow_kwargs).create_network() elif model_name == 'linear': model = nn.Sequential() model.add_module("conv_classifier", nn.Conv2d(n_chans, n_classes, (600, 1))) model.add_module('softmax', nn.LogSoftmax()) model.add_module('squeeze', Expression(lambda x: x.squeeze(3))) elif model_name == '3path': virtual_chan_1x1_conv = True mean_across_features = False drop_prob = 0.5 n_start_filters = 10 early_bnorm = False n_classifier_filters = 100 later_kernel_len = 5 extra_conv_stride = 4 # dont forget to reset n_preds_per_blabla model = create_multi_start_path_net( in_chans=n_chans, virtual_chan_1x1_conv=virtual_chan_1x1_conv, n_start_filters=n_start_filters, early_bnorm=early_bnorm, later_kernel_len=later_kernel_len, extra_conv_stride=extra_conv_stride, mean_across_features=mean_across_features, n_classifier_filters=n_classifier_filters, drop_prob=drop_prob) else: assert False, "unknown model name {:s}".format(model_name) if not model_name == '3path': to_dense_prediction_model(model) log.info("Model:\n{:s}".format(str(model))) time_cut_off_sec = np.inf start_time = time.time() # fix input time length in case of smac models if fix_input_length_for_smac: assert ('smac' in model_name) and (input_time_length == 12000) if cuda: model.cuda() test_input = np_to_var( np.ones((2, n_chans, input_time_length, 1), dtype=np.float32)) if cuda: test_input = test_input.cuda() try: out = model(test_input) except: raise ValueError("Model receptive field too large...") n_preds_per_input = out.cpu().data.numpy().shape[2] n_receptive_field = input_time_length - n_preds_per_input input_time_length = 2 * n_receptive_field exp = common.run_exp( max_recording_mins, n_recordings, sec_to_cut_at_start, sec_to_cut_at_end, duration_recording_mins, max_abs_val, clip_before_resample, sampling_freq, divisor, n_folds, i_test_fold, shuffle, merge_train_valid, model, input_time_length, optimizer, learning_rate, weight_decay, scheduler, model_constraint, batch_size, max_epochs, only_return_exp, time_cut_off_sec, start_time, test_on_eval, test_recording_mins, sensor_types, log_dir, np_th_seed, ) return exp
def runModel(mode): cudnn.benchmark = True start = time.time() #mode = str(sys.argv[1]) #X,y,test_X,test_y = loadSubNormData(mode='all') #X,y,test_X,test_y = loadNEDCdata(mode=mode) #data = np.load('sessionsData/data%s-sessions.npy'%mode[:3]) #labels = np.load('sessionsData/labels%s-sessions.npy'%mode[:3]) data = np.load('data%s.npy' % mode[:3]) labels = np.load('labels%s.npy' % mode[:3]) X, y, test_X, test_y = splitDataRandom_Loaded(data, labels, mode) print('Mode - %s Total n: %d, Test n: %d' % (mode, len(y) + len(test_y), len(test_y))) #return 0 #X = addDataNoise(X,band=[1,4]) #test_X = addDataNoise(test_X,band=[1,4]) max_shape = np.max([list(x.shape) for x in X], axis=0) assert max_shape[1] == int(config.duration_recording_mins * config.sampling_freq * 60) n_classes = 2 n_recordings = None # set to an integer, if you want to restrict the set size sensor_types = ["EEG"] n_chans = 19 #21 max_recording_mins = 35 # exclude larger recordings from training set sec_to_cut = 60 # cut away at start of each recording duration_recording_mins = 5 #20 # how many minutes to use per recording test_recording_mins = 5 #20 max_abs_val = 800 # for clipping sampling_freq = 100 divisor = 10 # divide signal by this test_on_eval = True # teston evaluation set or on training set # in case of test on eval, n_folds and i_testfold determine # validation fold in training set for training until first stop n_folds = 10 i_test_fold = 9 shuffle = True model_name = 'linear' #'deep'#'shallow' 'linear' n_start_chans = 25 n_chan_factor = 2 # relevant for deep model only input_time_length = 6000 final_conv_length = 1 model_constraint = 'defaultnorm' init_lr = 1e-3 batch_size = 64 max_epochs = 35 # until first stop, the continue train on train+valid cuda = True # False if model_name == 'shallow': model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, final_conv_length=final_conv_length).create_network() elif model_name == 'deep': model = Deep4Net(n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, stride_before_pool=True).create_network() elif (model_name == 'deep_smac'): if model_name == 'deep_smac': do_batch_norm = False else: assert model_name == 'deep_smac_bnorm' do_batch_norm = True double_time_convs = False drop_prob = 0.244445 filter_length_2 = 12 filter_length_3 = 14 filter_length_4 = 12 filter_time_length = 21 final_conv_length = 1 first_nonlin = elu first_pool_mode = 'mean' first_pool_nonlin = identity later_nonlin = elu later_pool_mode = 'mean' later_pool_nonlin = identity n_filters_factor = 1.679066 n_filters_start = 32 pool_time_length = 1 pool_time_stride = 2 split_first_layer = True n_chan_factor = n_filters_factor n_start_chans = n_filters_start model = Deep4Net(n_chans, n_classes, n_filters_time=n_start_chans, n_filters_spat=n_start_chans, input_time_length=input_time_length, n_filters_2=int(n_start_chans * n_chan_factor), n_filters_3=int(n_start_chans * (n_chan_factor**2.0)), n_filters_4=int(n_start_chans * (n_chan_factor**3.0)), final_conv_length=final_conv_length, batch_norm=do_batch_norm, double_time_convs=double_time_convs, drop_prob=drop_prob, filter_length_2=filter_length_2, filter_length_3=filter_length_3, filter_length_4=filter_length_4, filter_time_length=filter_time_length, first_nonlin=first_nonlin, first_pool_mode=first_pool_mode, first_pool_nonlin=first_pool_nonlin, later_nonlin=later_nonlin, later_pool_mode=later_pool_mode, later_pool_nonlin=later_pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, stride_before_pool=True).create_network() elif model_name == 'shallow_smac': conv_nonlin = identity do_batch_norm = True drop_prob = 0.328794 filter_time_length = 56 final_conv_length = 22 n_filters_spat = 73 n_filters_time = 24 pool_mode = 'max' pool_nonlin = identity pool_time_length = 84 pool_time_stride = 3 split_first_layer = True model = ShallowFBCSPNet( in_chans=n_chans, n_classes=n_classes, n_filters_time=n_filters_time, n_filters_spat=n_filters_spat, input_time_length=input_time_length, final_conv_length=final_conv_length, conv_nonlin=conv_nonlin, batch_norm=do_batch_norm, drop_prob=drop_prob, filter_time_length=filter_time_length, pool_mode=pool_mode, pool_nonlin=pool_nonlin, pool_time_length=pool_time_length, pool_time_stride=pool_time_stride, split_first_layer=split_first_layer, ).create_network() elif model_name == 'linear': model = nn.Sequential() model.add_module("conv_classifier", nn.Conv2d(n_chans, n_classes, (600, 1))) model.add_module('softmax', nn.LogSoftmax(dim=1)) model.add_module('squeeze', Expression(lambda x: x.squeeze(3))) else: assert False, "unknown model name {:s}".format(model_name) to_dense_prediction_model(model) if config.cuda: model.cuda() test_input = np_to_var( np.ones((2, config.n_chans, config.input_time_length, 1), dtype=np.float32)) if config.cuda: test_input = test_input.cuda() out = model(test_input) n_preds_per_input = out.cpu().data.numpy().shape[2] iterator = CropsFromTrialsIterator( batch_size=config.batch_size, input_time_length=config.input_time_length, n_preds_per_input=n_preds_per_input) #model.add_module('softmax', nn.LogSoftmax(dim=1)) model.eval() mode[2] = str(mode[2]) mode[3] = str(mode[3]) modelName = '-'.join(mode[:4]) #params = th.load('sessionsData/%sModel%s-sessions.pt'%(modelName,mode[4])) #params = th.load('%sModel%s.pt'%(modelName,mode[4])) params = th.load('linear/%sModel%s.pt' % (modelName, mode[4])) model.load_state_dict(params) if config.test_on_eval: #test_X, test_y = test_dataset.load() #test_X, test_y = loadNEDCdata(mode='eval') max_shape = np.max([list(x.shape) for x in test_X], axis=0) assert max_shape[1] == int(config.test_recording_mins * config.sampling_freq * 60) if not config.test_on_eval: splitter = TrainValidTestSplitter(config.n_folds, config.i_test_fold, shuffle=config.shuffle) train_set, valid_set, test_set = splitter.split(X, y) else: splitter = TrainValidSplitter(config.n_folds, i_valid_fold=config.i_test_fold, shuffle=config.shuffle) train_set, valid_set = splitter.split(X, y) test_set = SignalAndTarget(test_X, test_y) del test_X, test_y del X, y # shouldn't be necessary, but just to make sure datasets = OrderedDict( (('train', train_set), ('valid', valid_set), ('test', test_set))) for setname in ('train', 'valid', 'test'): #setname = 'test' #print("Compute predictions for {:s}...".format(setname)) dataset = datasets[setname] if config.cuda: preds_per_batch = [ var_to_np(model(np_to_var(b[0]).cuda())) for b in iterator.get_batches(dataset, shuffle=False) ] else: preds_per_batch = [ var_to_np(model(np_to_var(b[0]))) for b in iterator.get_batches(dataset, shuffle=False) ] preds_per_trial = compute_preds_per_trial( preds_per_batch, dataset, input_time_length=iterator.input_time_length, n_stride=iterator.n_preds_per_input) mean_preds_per_trial = [ np.mean(preds, axis=(0, 2)) for preds in preds_per_trial ] mean_preds_per_trial = np.array(mean_preds_per_trial) all_pred_labels = np.argmax(mean_preds_per_trial, axis=1).squeeze() all_target_labels = dataset.y acc_per_class = [] for i_class in range(n_classes): mask = all_target_labels == i_class acc = np.mean(all_pred_labels[mask] == all_target_labels[mask]) acc_per_class.append(acc) misclass = 1 - np.mean(acc_per_class) #print('Acc:{}, Class 0:{}, Class 1:{}'.format(np.mean(acc_per_class),acc_per_class[0],acc_per_class[1])) if setname == 'test': testResult = np.mean(acc_per_class) return testResult
def run_exp( debug, subject_id, max_epochs, n_sensors, final_hz, half_before, start_ms, stop_ms, model, weight_decay, final_fft, add_bnorm, act_norm, ): model_name = model del model assert final_hz in [64, 256] car = not debug train_inputs, test_inputs = load_train_test( subject_id, car, n_sensors, final_hz, start_ms, stop_ms, half_before, only_load_given_sensors=debug, ) cuda = True if cuda: train_inputs = [i.cuda() for i in train_inputs] test_inputs = [i.cuda() for i in test_inputs] from braindecode.datautil.signal_target import SignalAndTarget sets = [] for inputs in (train_inputs, test_inputs): X = np.concatenate([var_to_np(ins) for ins in inputs]).astype( np.float32 ) y = np.concatenate( [np.ones(len(ins)) * i_class for i_class, ins in enumerate(inputs)] ) y = y.astype(np.int64) set = SignalAndTarget(X, y) sets.append(set) train_set = sets[0] valid_set = sets[1] from braindecode.models.shallow_fbcsp import ShallowFBCSPNet from braindecode.models.deep4 import Deep4Net from torch import nn from braindecode.torch_ext.util import set_random_seeds set_random_seeds(2019011641, cuda) n_chans = train_inputs[0].shape[1] n_time = train_inputs[0].shape[2] n_classes = 2 input_time_length=train_set.X.shape[2] if model_name == 'shallow': # final_conv_length = auto ensures we only get a single output in the time dimension model = ShallowFBCSPNet(in_chans=n_chans, n_classes=n_classes, input_time_length=input_time_length, final_conv_length='auto') elif model_name == 'deep': model = Deep4Net(n_chans, n_classes, input_time_length=train_set.X.shape[2], pool_time_length=2, pool_time_stride=2, final_conv_length='auto') elif model_name == 'invertible': model = InvertibleModel(n_chans, n_time, final_fft=final_fft, add_bnorm=add_bnorm) elif model_name == 'deep_invertible': n_chan_pad = 0 filter_length_time = 11 model = deep_invertible( n_chans, input_time_length, n_chan_pad, filter_length_time) model.add_module("select_dims", Expression(lambda x: x[:, :2, 0])) model.add_module("softmax", nn.LogSoftmax(dim=1)) model = WrappedModel(model) ## set scale if act_norm: model.cuda() for module in model.network.modules(): if hasattr(module, 'log_factor'): module._forward_hooks.clear() module.register_forward_hook(scale_to_unit_var) model.network(train_inputs[0].cuda()); for module in model.network.modules(): if hasattr(module, 'log_factor'): module._forward_hooks.clear() else: assert False if cuda: model.cuda() from braindecode.torch_ext.optimizers import AdamW import torch.nn.functional as F if model_name == 'shallow': assert weight_decay == 'hardcoded' optimizer = AdamW(model.parameters(), lr=0.0625 * 0.01, weight_decay=0) elif model_name == 'deep': assert weight_decay == 'hardcoded' optimizer = AdamW(model.parameters(), lr=1 * 0.01, weight_decay=0.5 * 0.001) # these are good values for the deep model elif model_name == 'invertible': optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=weight_decay) elif model_name == 'deep_invertible': optimizer = AdamW(model.parameters(), lr=1 * 0.001, weight_decay=weight_decay) else: assert False model.compile(loss=F.nll_loss, optimizer=optimizer, iterator_seed=1, ) model.fit(train_set.X, train_set.y, epochs=max_epochs, batch_size=64, scheduler='cosine', validation_data=(valid_set.X, valid_set.y), ) return model.epochs_df, model.network