eval_base.py

#!/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()