--- a +++ b/eval_base.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python +# coding: utf-8 +'''Subject-independent model evaluator. +''' +import argparse +import json +import logging +import sys +from os.path import join as pjoin + +import h5py +import numpy as np +import torch +import torch.nn.functional as F +from braindecode.models.deep4 import Deep4Net +from braindecode.torch_ext.optimizers import AdamW +from braindecode.torch_ext.util import set_random_seeds + +logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s', + level=logging.INFO, stream=sys.stdout) + +parser = argparse.ArgumentParser( + description='Subject independent model evaluator.') +parser.add_argument('datapath', type=str, help='Path to KU data') +parser.add_argument('modelpath', type=str, help='Path to base model') +parser.add_argument('outpath', type=str, help='Path to output') +parser.add_argument('-gpu', type=int, help='The gpu device to use', default=0) + +args = parser.parse_args() +datapath = args.datapath +outpath = args.outpath +modelpath = args.modelpath +dfile = h5py.File(datapath, 'r') +torch.cuda.set_device(args.gpu) +set_random_seeds(seed=20200205, cuda=True) +BATCH_SIZE = 16 + +# Randomly shuffled subject. +subjs = [35, 47, 46, 37, 13, 27, 12, 32, 53, 54, 4, 40, 19, 41, 18, 42, 34, 7, + 49, 9, 5, 48, 29, 15, 21, 17, 31, 45, 1, 38, 51, 8, 11, 16, 28, 44, 24, + 52, 3, 26, 39, 50, 6, 23, 2, 14, 25, 20, 10, 33, 22, 43, 36, 30] + + +# Get data from single subject. +def get_data(subj): + dpath = '/s' + str(subj) + X = dfile[pjoin(dpath, 'X')] + Y = dfile[pjoin(dpath, 'Y')] + return X[:], Y[:] + + +X, Y = get_data(subjs[0]) +n_classes = 2 +in_chans = X.shape[1] +# final_conv_length = auto ensures we only get a single output in the time dimension +model = Deep4Net(in_chans=in_chans, n_classes=n_classes, + input_time_length=X.shape[2], + final_conv_length='auto').cuda() + +# Dummy train data to set up the model. +X_train = np.zeros(X[:2].shape).astype(np.float32) +Y_train = np.zeros(Y[:2].shape).astype(np.int64) + + +def reset_model(checkpoint): + # Load the state dict of the model. + model.network.load_state_dict(checkpoint['model_state_dict']) + + # Only optimize parameters that requires gradient. + optimizer = AdamW(filter(lambda p: p.requires_grad, model.network.parameters()), + lr=1*0.01, weight_decay=0.5*0.001) + model.compile(loss=F.nll_loss, optimizer=optimizer, + iterator_seed=20200205, ) + + +for fold, subj in enumerate(subjs): + suffix = '_s' + str(subj) + '_f' + str(fold) + checkpoint = torch.load(pjoin(modelpath, 'model_f' + str(fold) + '.pt'), + map_location='cuda:' + str(args.gpu)) + + # Set up the model. + reset_model(checkpoint) + model.fit(X_train, Y_train, 0, BATCH_SIZE) + + X, Y = get_data(subj) + X_test, Y_test = X[300:], Y[300:] + test_loss = model.evaluate(X_test, Y_test) + with open(pjoin(outpath, 'test_base' + suffix + '.json'), 'w') as f: + json.dump(test_loss, f) + +dfile.close()