def evaluate_test_set_performance(model_dir):
"""Measures the test set performance of the model under the specified model directory.
:param model_dir: directory containing the model state dictionaries for each fold and the model
configuration (including the population graph parameterisation)
:return: the test set performance for each fold.
"""
with open(os.path.join(model_dir, 'config.yaml')) as file:
cfg = yaml.full_load(file)
graph_name = cfg['graph_name']['value']
conv_type = cfg['model']['value']
n_conv_layers = cfg['n_conv_layers']['value']
layer_sizes = ast.literal_eval(cfg['layer_sizes']['value'])
dropout_p = cfg['dropout']['value']
similarity_feature_set = [Phenotype(i) for i in ast.literal_eval(cfg['similarity']['value'])[0]]
similarity_threshold = ast.literal_eval(cfg['similarity']['value'])[1]
if graph_name not in GRAPH_NAMES:
graph_construct.construct_population_graph(similarity_feature_set=similarity_feature_set,
similarity_threshold=similarity_threshold,
functional=False,
structural=True,
euler=True)
graph = graph_construct.load_population_graph(graph_root, graph_name)
folds = brain_gnn_train.get_cv_subject_split(graph, n_folds=5)
results = {}
for i, fold in enumerate(folds):
brain_gnn_train.set_training_masks(graph, *fold)
graph_transform.graph_feature_transform(graph)
if ConvTypes(conv_type) == ConvTypes.GCN:
model = BrainGCN(graph.num_node_features, n_conv_layers, layer_sizes, dropout_p)
else:
model = BrainGAT(graph.num_node_features, n_conv_layers, layer_sizes, dropout_p)
model.load_state_dict(torch.load(os.path.join(model_dir, 'fold-{}_state_dict.pt'.format(i))))
model = model.to(device)
model.eval()
data = graph.to(device)
model = model(data)
predicted = model[data.test_mask].cpu()
actual = graph.y[data.test_mask].cpu()
r2 = r2_score(actual.detach().numpy(), predicted.detach().numpy())
r = pearsonr(actual.detach().numpy().flatten(), predicted.detach().numpy().flatten())
results['fold_{}'.format(i)] = {'r': [x.item() for x in r], 'r2': r2.item()}
mse = mean_squared_error(actual.detach().numpy(), predicted.detach().numpy())
results=mse
break
return results
def train_with_cross_validation(conv_type,
graph,
device,
n_folds=10,
n_conv_layers=0,
layer_sizes=None,
epochs=350,
lr=0.005,
dropout_p=0,
weight_decay=1e-5,
log=True,
early_stopping=True,
patience=10,
delta=0.005):
"""
Trains the graph neural network on a population graph.
:param conv_type: convolution type, ConvType.GCN or ConvType.GAT
:param graph: the population graph
:param device: device on which the neural network will be trained. 'cpu' or 'cuda:x'
:param n_conv_layers: number of convolutional layers in GNN
:param layer_sizes: array of layer sizes, first n_conv_layers of which convolutional, the rest fully connected
:param epochs: number of epochs for which to train
:param lr: learning rate
:param dropout_p: dropout probability (probability of killing the unit)
:param weight_decay: weight decay
:param log: indicates whether to log results to wandb.
:param early_stopping: indicates whether to use early stopping.
:param patience: indicates how long to tolerate no improvement in MSE before early stopping.
:param delta: defines the threshold of how much the model has to improve in `patience` epochs.
:return: array of results as in train method, repeated for all folds.
"""
folds = get_cv_subject_split(graph, n_folds=n_folds)
results = []
run_name = None
wandb.init(project="brain-age-gnn",
config=hyperparameter_defaults,
reinit=True)
wandb.save("*.pt")
for i, fold in enumerate(folds):
set_training_masks(graph, *fold)
graph_transform.graph_feature_transform(graph)
_, fold_result = train(conv_type,
graph,
device,
n_conv_layers=n_conv_layers,
layer_sizes=layer_sizes,
epochs=epochs,
lr=lr,
dropout_p=dropout_p,
weight_decay=weight_decay,
log=log,
early_stopping=early_stopping,
patience=patience,
delta=delta,
cv=True,
fold=i,
run_name=run_name)
run_name = fold_result[0]
fold_scores = fold_result[1:]
results.append(fold_scores)
if np.mean([
F.mse_loss(predicted, actual).item()
for predicted, actual in results
],
axis=0) > 30:
break
# Add cross-validation summaries to the last fold
cv_mse = [
F.mse_loss(predicted, actual).item() for predicted, actual in results
]
cv_r2 = [
r2_score(actual.detach().numpy(),
predicted.detach().numpy()) for predicted, actual in results
]
cv_r = [
pearsonr(actual.detach().numpy().flatten(),
predicted.detach().numpy().flatten())
for predicted, actual in results
]
wandb.run.summary["cv_validation_average_mse"] = np.mean(cv_mse, axis=0)
wandb.run.summary["cv_validation_average_r2"] = np.mean(cv_r2, axis=0)
wandb.run.summary["cv_validation_average_r"] = np.mean(cv_r, axis=0)
return results
def label_permutation_test(model_dir):
"""Permutation test measuring the performance of the model when the labels are shuffled.
:param model_dir: directory containing the model state dictionaries for each fold and the model
configuration (including the population graph parameterisation)
:return: the test set performance for each permutation.
"""
with open(os.path.join(model_dir, 'config.yaml')) as file:
cfg = yaml.full_load(file)
graph_name = cfg['graph_name']['value']
conv_type = cfg['model']['value']
n_conv_layers = cfg['n_conv_layers']['value']
layer_sizes = ast.literal_eval(cfg['layer_sizes']['value'])
dropout_p = cfg['dropout']['value']
similarity_feature_set = [Phenotype(i) for i in ast.literal_eval(cfg['similarity']['value'])[0]]
similarity_threshold = ast.literal_eval(cfg['similarity']['value'])[1]
if graph_name not in GRAPH_NAMES:
graph_construct.construct_population_graph(similarity_feature_set=similarity_feature_set,
similarity_threshold=similarity_threshold,
functional=False,
structural=True,
euler=True)
graph = graph_construct.load_population_graph(graph_root, graph_name)
folds = brain_gnn_train.get_cv_subject_split(graph, n_folds=5)
fold = folds[0]
brain_gnn_train.set_training_masks(graph, *fold)
graph_transform.graph_feature_transform(graph)
rs = []
r2s = []
mses = []
for i in range(1000):
graph.to('cpu')
permute_population_graph_labels(graph, i)
if ConvTypes(conv_type) == ConvTypes.GCN:
model = BrainGCN(graph.num_node_features, n_conv_layers, layer_sizes, dropout_p)
else:
model = BrainGAT(graph.num_node_features, n_conv_layers, layer_sizes, dropout_p)
model.load_state_dict(torch.load(os.path.join(model_dir, 'fold-{}_state_dict.pt'.format(0))))
model = model.to(device)
data = graph.to(device)
model.eval()
model = model(data)
predicted = model[data.test_mask].cpu()
actual = graph.y[data.test_mask].cpu()
r2 = r2_score(actual.detach().numpy(), predicted.detach().numpy())
r = pearsonr(actual.detach().numpy().flatten(), predicted.detach().numpy().flatten())
mse = mean_squared_error(actual.detach().numpy(), predicted.detach().numpy())
rs.append(r[0])
r2s.append(r2)
mses.append(mse)
print(r[0], r2, mse)
np.save(os.path.join('notebooks', 'permutations_{}_{}'.format(conv_type, 'r')), rs)
np.save(os.path.join('notebooks', 'permutations_{}_{}'.format(conv_type, 'r2')), r2s)
np.save(os.path.join('notebooks', 'permutations_{}_{}'.format(conv_type, 'mse')), mses)
return [rs, r2s]
def evaluate_noise_performance(model_dir, noise_type='node'):
"""Measures the test set performance of the model under the specified model directory when noise is added.
:param model_dir: directory containing the model state dictionaries for each fold and the model
configuration (including the population graph parameterisation)
:param noise_type: 'node', 'node_feature_permutation' or 'edge'.
:return: the dictionary of results under five different random seeds and increasing probabilities of added noise.
"""
with open(os.path.join(model_dir, 'config.yaml')) as file:
cfg = yaml.full_load(file)
graph_name = cfg['graph_name']['value']
conv_type = cfg['model']['value']
n_conv_layers = cfg['n_conv_layers']['value']
layer_sizes = ast.literal_eval(cfg['layer_sizes']['value'])
dropout_p = cfg['dropout']['value']
lr = cfg['learning_rate']['value']
weight_decay = cfg['weight_decay']['value']
similarity_feature_set = [Phenotype(i) for i in ast.literal_eval(cfg['similarity']['value'])[0]]
similarity_threshold = ast.literal_eval(cfg['similarity']['value'])[1]
if graph_name not in GRAPH_NAMES:
graph_construct.construct_population_graph(similarity_feature_set=similarity_feature_set,
similarity_threshold=similarity_threshold,
functional=False,
structural=True,
euler=True)
graph = graph_construct.load_population_graph(graph_root, graph_name)
folds = brain_gnn_train.get_cv_subject_split(graph, n_folds=5)
fold = folds[0]
results = {}
for i in range(1, 5):
brain_gnn_train.set_training_masks(graph, *fold)
results_fold = {}
for p in [0.01, 0.05, 0.1, 0.2, 0.3, 0.5, 0.8, 0.95]:
graph.to('cpu')
graph_transform.graph_feature_transform(graph)
if noise_type == 'node':
add_population_graph_noise(graph, p, random_state=i)
if noise_type == 'edge':
remove_population_graph_edges(graph, p, random_state=i)
if noise_type == 'node-feature-permutation':
permute_population_graph_features(graph, p, random_state=i)
data = graph.to(device)
epochs = 10000
model, _ = brain_gnn_train.train(conv_type, graph, device, n_conv_layers, layer_sizes, epochs, lr,
dropout_p, weight_decay, patience=100)
model.eval()
model = model(data)
predicted = model[data.test_mask].cpu()
actual = data.y[data.test_mask].cpu()
r2 = r2_score(actual.detach().numpy(), predicted.detach().numpy())
r = pearsonr(actual.detach().numpy().flatten(), predicted.detach().numpy().flatten())
results_fold['p={}_metric=r'.format(p)] = [x.item() for x in r][0]
wandb.run.summary['{}_{}_{}_p={}_metric=r'.format(conv_type, noise_type, i, p)] = [x.item() for x in r][0]
results_fold['p={}_metric=r2'.format(p)] = r2.item()
wandb.run.summary['{}_{}_{}_p={}_metric=r2'.format(conv_type, noise_type, i, p)] = r2.item()
gc.collect()
results['{}_{}_{}'.format(conv_type, noise_type, i)] = results_fold
return results
structural=structural,
euler=euler)
if args.model == 'gcn' or args.model == 'gat':
n_conv_layers = args.n_conv_layers
else:
n_conv_layers = 0
torch.manual_seed(99)
np.random.seed(0)
population_graph = graph_construct.load_population_graph(
graph_root, graph_name)
fold = brain_gnn_train.get_stratified_subject_split(population_graph)
brain_gnn_train.set_training_masks(population_graph, *fold)
graph_transform.graph_feature_transform(population_graph)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
brain_gnn_train.train_with_cross_validation(args.model,
population_graph,
device,
n_conv_layers=n_conv_layers,
layer_sizes=ast.literal_eval(
args.layer_sizes),
lr=args.learning_rate,
weight_decay=args.weight_decay,
dropout_p=args.dropout,
epochs=args.epochs,
n_folds=5,
patience=100)