# -*- coding: utf-8 -*-
import os
import gc
import time
import numpy as np
from collections import defaultdict
from keras import backend as K
from keras import optimizers
from utils import load_data, pickle_load, format_filename, write_log
from models import KGCN
from config import ModelConfig, PROCESSED_DATA_DIR, ENTITY_VOCAB_TEMPLATE, \
RELATION_VOCAB_TEMPLATE, ADJ_ENTITY_TEMPLATE, ADJ_RELATION_TEMPLATE, LOG_DIR, PERFORMANCE_LOG, \
DRUG_VOCAB_TEMPLATE
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
def get_optimizer(op_type, learning_rate):
if op_type == 'sgd':
return optimizers.SGD(learning_rate)
elif op_type == 'rmsprop':
return optimizers.RMSprop(learning_rate)
elif op_type == 'adagrad':
return optimizers.Adagrad(learning_rate)
elif op_type == 'adadelta':
return optimizers.Adadelta(learning_rate)
elif op_type == 'adam':
return optimizers.Adam(learning_rate, clipnorm=5)
else:
raise ValueError('Optimizer Not Understood: {}'.format(op_type))
def train(train_d,dev_d,test_d,kfold,dataset, neighbor_sample_size, embed_dim, n_depth, l2_weight, lr, optimizer_type,
batch_size, aggregator_type, n_epoch, callbacks_to_add=None, overwrite=True):
config = ModelConfig()
config.neighbor_sample_size = neighbor_sample_size
config.embed_dim = embed_dim
config.n_depth = n_depth
config.l2_weight = l2_weight
config.dataset=dataset
config.K_Fold=kfold
config.lr = lr
config.optimizer = get_optimizer(optimizer_type, lr)
config.batch_size = batch_size
config.aggregator_type = aggregator_type
config.n_epoch = n_epoch
config.callbacks_to_add = callbacks_to_add
config.drug_vocab_size = len(pickle_load(format_filename(PROCESSED_DATA_DIR,
DRUG_VOCAB_TEMPLATE,
dataset=dataset)))
config.entity_vocab_size = len(pickle_load(format_filename(PROCESSED_DATA_DIR,
ENTITY_VOCAB_TEMPLATE,
dataset=dataset)))
config.relation_vocab_size = len(pickle_load(format_filename(PROCESSED_DATA_DIR,
RELATION_VOCAB_TEMPLATE,
dataset=dataset)))
config.adj_entity = np.load(format_filename(PROCESSED_DATA_DIR, ADJ_ENTITY_TEMPLATE,
dataset=dataset))
config.adj_relation = np.load(format_filename(PROCESSED_DATA_DIR, ADJ_RELATION_TEMPLATE,
dataset=dataset))
config.exp_name = f'kgcn_{dataset}_neigh_{neighbor_sample_size}_embed_{embed_dim}_depth_' \
f'{n_depth}_agg_{aggregator_type}_optimizer_{optimizer_type}_lr_{lr}_' \
f'batch_size_{batch_size}_epoch_{n_epoch}'
callback_str = '_' + '_'.join(config.callbacks_to_add)
callback_str = callback_str.replace('_modelcheckpoint', '').replace('_earlystopping', '')#去掉了这两种方式使用swa得方式平均
config.exp_name += callback_str
train_log = {'exp_name': config.exp_name, 'batch_size': batch_size, 'optimizer': optimizer_type,
'epoch': n_epoch, 'learning_rate': lr}
print('Logging Info - Experiment: %s' % config.exp_name)
model_save_path = os.path.join(config.checkpoint_dir, '{}.hdf5'.format(config.exp_name))
model = KGCN(config)
train_data=np.array(train_d)
valid_data=np.array(dev_d)
test_data=np.array(test_d)
if not os.path.exists(model_save_path) or overwrite:
start_time = time.time()
model.fit(x_train=[train_data[:, :1], train_data[:, 1:2]], y_train=train_data[:, 2:3],
x_valid=[valid_data[:, :1], valid_data[:, 1:2]], y_valid=valid_data[:, 2:3])
elapsed_time = time.time() - start_time
print('Logging Info - Training time: %s' % time.strftime("%H:%M:%S",
time.gmtime(elapsed_time)))
train_log['train_time'] = time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
print('Logging Info - Evaluate over valid data:')
model.load_best_model()
auc, acc, f1,aupr = model.score(x=[valid_data[:, :1], valid_data[:, 1:2]], y=valid_data[:, 2:3])
print(f'Logging Info - dev_auc: {auc}, dev_acc: {acc}, dev_f1: {f1}, dev_aupr: {aupr}'
)
train_log['dev_auc'] = auc
train_log['dev_acc'] = acc
train_log['dev_f1'] = f1
train_log['dev_aupr']=aupr
train_log['k_fold']=kfold
train_log['dataset']=dataset
train_log['aggregate_type']=config.aggregator_type
if 'swa' in config.callbacks_to_add:
model.load_swa_model()
print('Logging Info - Evaluate over valid data based on swa model:')
auc, acc, f1,aupr = model.score(x=[valid_data[:, :1], valid_data[:, 1:2]], y=valid_data[:, 2:3])
train_log['swa_dev_auc'] = auc
train_log['swa_dev_acc'] = acc
train_log['swa_dev_f1'] = f1
train_log['swa_dev_aupr']=aupr
print(f'Logging Info - swa_dev_auc: {auc}, swa_dev_acc: {acc}, swa_dev_f1: {f1}, swa_dev_aupr: {aupr}') #修改输出指标
print('Logging Info - Evaluate over test data:')
model.load_best_model()
auc, acc, f1, aupr = model.score(x=[test_data[:, :1], test_data[:, 1:2]], y=test_data[:, 2:3])
train_log['test_auc'] = auc
train_log['test_acc'] = acc
train_log['test_f1'] = f1
train_log['test_aupr'] =aupr
print(f'Logging Info - test_auc: {auc}, test_acc: {acc}, test_f1: {f1}, test_aupr: {aupr}')
if 'swa' in config.callbacks_to_add:
model.load_swa_model()
print('Logging Info - Evaluate over test data based on swa model:')
auc, acc, f1,aupr = model.score(x=[test_data[:, :1], test_data[:, 1:2]], y=test_data[:, 2:3])
train_log['swa_test_auc'] = auc
train_log['swa_test_acc'] = acc
train_log['swa_test_f1'] = f1
train_log['swa_test_aupr'] = aupr
print(f'Logging Info - swa_test_auc: {auc}, swa_test_acc: {acc}, swa_test_f1: {f1}, swa_test_aupr: {aupr}')
train_log['timestamp'] = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
write_log(format_filename(LOG_DIR, PERFORMANCE_LOG), log=train_log, mode='a')
del model
gc.collect()
K.clear_session()
return train_log
Datasets
Models
