def batch_dataset_generator(gen, args, is_testing=False):
grid_size = subgrid_gen.grid_size(args.grid_config)
channel_size = subgrid_gen.num_channels(args.grid_config)
dataset = tf.data.Dataset.from_generator(
gen,
output_types=(tf.string, tf.float32, tf.float32),
output_shapes=((), (grid_size, grid_size, grid_size, channel_size),
(1, )))
# Shuffle dataset
if not is_testing:
if args.shuffle:
dataset = dataset.repeat(count=None)
else:
dataset = dataset.apply(
tf.contrib.data.shuffle_and_repeat(buffer_size=1000))
dataset = dataset.batch(args.batch_size)
dataset = dataset.prefetch(8)
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()
return dataset, next_element
def __channel_size(args):
size = subgrid_gen.num_channels(args.grid_config)
if args.add_flag:
size += 1
return size
def train_model(sess, args):
# tf Graph input
# Subgrid maps for each residue in a protein
logging.debug('Create input placeholder...')
grid_size = subgrid_gen.grid_size(args.grid_config)
channel_size = subgrid_gen.num_channels(args.grid_config)
feature_placeholder = tf.placeholder(
tf.float32, [None, grid_size, grid_size, grid_size, channel_size],
name='main_input')
label_placeholder = tf.placeholder(tf.float32, [None, 1], 'label')
# Placeholder for model parameters
training_placeholder = tf.placeholder(tf.bool,
shape=[],
name='is_training')
conv_drop_rate_placeholder = tf.placeholder(tf.float32,
name='conv_drop_rate')
fc_drop_rate_placeholder = tf.placeholder(tf.float32, name='fc_drop_rate')
top_nn_drop_rate_placeholder = tf.placeholder(tf.float32,
name='top_nn_drop_rate')
# Define loss and optimizer
logging.debug('Define loss and optimizer...')
predict_op, loss_op = conv_model(feature_placeholder, label_placeholder,
training_placeholder,
conv_drop_rate_placeholder,
fc_drop_rate_placeholder,
top_nn_drop_rate_placeholder, args)
logging.debug('Generate training ops...')
train_op = model.training(loss_op, args.learning_rate)
# Initialize the variables (i.e. assign their default value)
logging.debug('Initializing global variables...')
init = tf.global_variables_initializer()
# Create saver and summaries.
logging.debug('Initializing saver...')
saver = tf.train.Saver(max_to_keep=100000)
logging.debug('Finished initializing saver...')
def __loop(generator, mode, num_iters):
tf_dataset, next_element = batch_dataset_generator(
generator, args, is_testing=(mode == 'test'))
structs, losses, preds, labels = [], [], [], []
epoch_loss = 0
progress_format = mode + ' loss: {:6.6f}'
# Loop over all batches (one batch is all feature for 1 protein)
num_batches = int(math.ceil(float(num_iters) / args.batch_size))
#print('Running {:} -> {:} iters in {:} batches (batch size: {:})'.format(
# mode, num_iters, num_batches, args.batch_size))
with tqdm.tqdm(total=num_batches, desc=progress_format.format(0)) as t:
for i in range(num_batches):
try:
struct_, feature_, label_ = sess.run(next_element)
_, pred, loss = sess.run(
[train_op, predict_op, loss_op],
feed_dict={
feature_placeholder:
feature_,
label_placeholder:
label_,
training_placeholder: (mode == 'train'),
conv_drop_rate_placeholder:
args.conv_drop_rate if mode == 'train' else 0.0,
fc_drop_rate_placeholder:
args.fc_drop_rate if mode == 'train' else 0.0,
top_nn_drop_rate_placeholder:
args.top_nn_drop_rate if mode == 'train' else 0.0
})
epoch_loss += (np.mean(loss) - epoch_loss) / (i + 1)
structs.extend(struct_.astype(str))
losses.append(loss)
preds.extend(pred)
labels.extend(label_)
t.set_description(progress_format.format(epoch_loss))
t.update(1)
except (tf.errors.OutOfRangeError, StopIteration):
logging.info("\nEnd of dataset at iteration {:}".format(i))
break
def __concatenate(array):
try:
array = np.concatenate(array)
return array
except:
return array
structs = __concatenate(structs)
preds = __concatenate(preds)
labels = __concatenate(labels)
losses = __concatenate(losses)
return structs, preds, labels, losses, epoch_loss
# Run the initializer
logging.debug('Running initializer...')
sess.run(init)
logging.debug('Finished running initializer...')
##### Training + validation
if not args.test_only:
prev_val_loss, best_val_loss = float("inf"), float("inf")
if (args.max_targets_train == None) and (args.max_decoys_train
== None):
train_num_structs = args.train_sharded.get_num_structures(
['ensemble', 'subunit'])
elif (args.max_targets_train == None):
train_num_structs = args.train_sharded.get_num_keyed(
) * args.max_decoys_train
elif (args.max_decoys_train == None):
assert False
else:
train_num_structs = args.max_targets_train * args.max_decoys_train
if (args.max_targets_val == None) and (args.max_decoys_val == None):
val_num_structs = args.val_sharded.get_num_structures(
['ensemble', 'subunit'])
elif (args.max_targets_val == None):
val_num_structs = args.val_sharded.get_num_keyed(
) * args.max_decoys_val
elif (args.max_decoys_val == None):
assert False
else:
val_num_structs = args.max_targets_val * args.max_decoys_val
train_num_structs *= args.repeat_gen
#val_num_structs *= args.repeat_gen
logging.info(
"Start training with {:} structs for train and {:} structs for val per epoch"
.format(train_num_structs, val_num_structs))
def _save():
ckpt = saver.save(sess,
os.path.join(args.output_dir, 'model-ckpt'),
global_step=epoch)
return ckpt
run_info_filename = os.path.join(args.output_dir, 'run_info.json')
run_info = {}
def __update_and_write_run_info(key, val):
run_info[key] = val
with open(run_info_filename, 'w') as f:
json.dump(run_info, f, indent=4)
per_epoch_val_losses = []
for epoch in range(1, args.num_epochs + 1):
random_seed = args.random_seed #random.randint(1, 10e6)
logging.info('Epoch {:} - random_seed: {:}'.format(
epoch, args.random_seed))
logging.debug('Creating train generator...')
train_generator_callable = functools.partial(
feature_psp.dataset_generator,
args.train_sharded,
args.grid_config,
score_type=args.score_type,
shuffle=args.shuffle,
repeat=args.repeat_gen,
max_targets=args.max_targets_train,
max_decoys=args.max_decoys_train,
max_dist_threshold=300.0,
random_seed=random_seed)
logging.debug('Creating val generator...')
val_generator_callable = functools.partial(
feature_psp.dataset_generator,
args.val_sharded,
args.grid_config,
score_type=args.score_type,
shuffle=args.shuffle,
repeat=1, #*args.repeat_gen,
max_targets=args.max_targets_val,
max_decoys=args.max_decoys_val,
max_dist_threshold=300.0,
random_seed=random_seed)
# Training
train_structs, train_preds, train_labels, _, curr_train_loss = __loop(
train_generator_callable, 'train', num_iters=train_num_structs)
# Validation
val_structs, val_preds, val_labels, _, curr_val_loss = __loop(
val_generator_callable, 'val', num_iters=val_num_structs)
per_epoch_val_losses.append(curr_val_loss)
__update_and_write_run_info('val_losses', per_epoch_val_losses)
if args.use_best or args.early_stopping:
if curr_val_loss < best_val_loss:
# Found new best epoch.
best_val_loss = curr_val_loss
ckpt = _save()
__update_and_write_run_info('val_best_loss', best_val_loss)
__update_and_write_run_info('best_ckpt', ckpt)
logging.info("New best {:}".format(ckpt))
if (epoch == args.num_epochs - 1 and not args.use_best):
# At end and just using final checkpoint.
ckpt = _save()
__update_and_write_run_info('best_ckpt', ckpt)
logging.info("Last checkpoint {:}".format(ckpt))
if args.save_all_ckpts:
# Save at every checkpoint
ckpt = _save()
logging.info("Saving checkpoint {:}".format(ckpt))
## Save last train and val results
train_df = pd.DataFrame(
np.array([train_structs, train_labels, train_preds]).T,
columns=['structure', 'true', 'pred'],
)
train_df['target'] = train_df.structure.apply(
lambda x: psp_util.get_target_name(x))
train_df.to_pickle(
os.path.join(args.output_dir,
'train_result-{:}.pkl'.format(epoch)))
__stats('Train Epoch {:}'.format(epoch), train_df)
val_df = pd.DataFrame(
np.array([val_structs, val_labels, val_preds]).T,
columns=['structure', 'true', 'pred'],
)
val_df['target'] = val_df.structure.apply(
lambda x: psp_util.get_target_name(x))
val_df.to_pickle(
os.path.join(args.output_dir,
'val_result-{:}.pkl'.format(epoch)))
__stats('Val Epoch {:}'.format(epoch), val_df)
if args.early_stopping and curr_val_loss >= prev_val_loss:
logging.info("Validation loss stopped decreasing, stopping...")
break
else:
prev_val_loss = curr_val_loss
logging.info("Finished training")
##### Testing
logging.debug("Run testing")
if not args.test_only:
to_use = run_info['best_ckpt'] if args.use_best else ckpt
else:
with open(os.path.join(args.model_dir, 'run_info.json')) as f:
run_info = json.load(f)
to_use = run_info['best_ckpt']
saver = tf.train.import_meta_graph(to_use + '.meta')
logging.info("Using {:} for testing".format(to_use))
saver.restore(sess, to_use)
test_generator_callable = functools.partial(
feature_psp.dataset_generator,
args.test_sharded,
args.grid_config,
score_type=args.score_type,
shuffle=args.shuffle,
repeat=1,
max_targets=args.max_targets_test,
max_decoys=args.max_decoys_test,
max_dist_threshold=None,
random_seed=args.random_seed)
if (args.max_targets_test == None) and (args.max_decoys_test == None):
test_num_structs = args.test_sharded.get_num_structures(
['ensemble', 'subunit'])
elif (args.max_targets_test == None):
test_num_structs = args.test_sharded.get_num_keyed(
) * args.max_decoys_test
elif (args.max_decoys_test == None):
assert False
else:
test_num_structs = args.max_targets_test * args.max_decoys_test
logging.info("Start testing with {:} structs".format(test_num_structs))
test_structs, test_preds, test_labels, _, test_loss = __loop(
test_generator_callable, 'test', num_iters=test_num_structs)
logging.info("Finished testing")
test_df = pd.DataFrame(
np.array([test_structs, test_labels, test_preds]).T,
columns=['structure', 'true', 'pred'],
)
test_df.to_pickle(os.path.join(args.output_dir, 'test_result.pkl'))
test_df['target'] = test_df.structure.apply(
lambda x: psp_util.get_target_name(x))
test_df.to_pickle(os.path.join(args.output_dir, 'test_result.pkl'))
__stats('Test', test_df)