"""
Documentation
"""
# Python Modules
import argparse
import logging
import os
import shutil
from typing import Any, Dict, Optional
# 3rd Party Modules
import numpy as np
import tensorflow as tf
# Project Modules
import deletor.tfutils as tfutils
from examples.utils import train, evaluate
from deletor.metrics import NormalizedDiscountedCumulativeGain
from deletor.models.mlp import ModelParameter, SimpleScoringNetwork
from examples.pipeline import load_dataset, truncate_document_list, is_valid_query, \
shuffle_documents, make_padded_shapes, make_padding_values
tfutils.grow_memory()
from deletor.losses import ApproximateNormalizedDiscountedCumulativeGain
np.set_printoptions(precision=6, suppress=True, edgeitems=10, linewidth=10000)
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)-15s [%(name)s]:%(lineno)d %(levelname)s %(message)s'
)
log = logging.getLogger('mlp/mltr30k')
AUTOTUNE = tf.data.experimental.AUTOTUNE
[docs]def prepare_data(
args: argparse.Namespace,
list_size: Optional[int],
train_bsz: int,
eval_bsz: int,
drop_remainder: bool = False
):
train_data = load_dataset(args.train_file, args.scaler)
valid_data = load_dataset(args.valid_file, args.scaler)
test_data = load_dataset(args.test_file, args.scaler)
train_data = train_data.filter(is_valid_query)
valid_data = valid_data.filter(is_valid_query)
test_data = test_data.filter(is_valid_query)
if list_size:
train_data = train_data.map(lambda x, y: truncate_document_list(x, y, list_size))
valid_data = valid_data.map(lambda x, y: truncate_document_list(x, y, list_size))
test_data = test_data.map(lambda x, y: truncate_document_list(x, y, list_size))
train_data = train_data.cache()
valid_data = valid_data.cache()
test_data = test_data.cache()
train_data = train_data.map(lambda x, y: shuffle_documents(x, y))
train_data = train_data.shuffle(1000, seed=args.random_seed, reshuffle_each_iteration=True)
padded_shapes = make_padded_shapes(list_size)
padding_values = make_padding_values()
train_data = train_data.padded_batch(train_bsz, padded_shapes, padding_values, drop_remainder)
valid_data = valid_data.padded_batch(eval_bsz, padded_shapes, padding_values, drop_remainder)
test_data = test_data.padded_batch(eval_bsz, padded_shapes, padding_values, drop_remainder)
train_data = train_data.prefetch(AUTOTUNE)
valid_data = valid_data.prefetch(AUTOTUNE)
test_data = test_data.prefetch(AUTOTUNE)
return train_data, valid_data, test_data
[docs]def setup_model(model_params: Dict[str, Any]):
model = SimpleScoringNetwork(model_params)
optimizer = tf.keras.optimizers.Adagrad(0.0075)
# optimizer = tf.keras.optimizers.SGD(0.0001, nesterov=True)
# optimizer = tf.keras.optimizers.RM
loss = ApproximateNormalizedDiscountedCumulativeGain(reduce=True)
metrics = [
NormalizedDiscountedCumulativeGain(k=1),
NormalizedDiscountedCumulativeGain(k=5),
NormalizedDiscountedCumulativeGain(k=10),
]
model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
return model
[docs]def main(args: argparse.Namespace):
list_size = None
train_batch_size = 32
eval_batch_size = 64
model_params = {
ModelParameter.N_UNITS: [256, 128, 64, 32, 16],
ModelParameter.DROPOUT_RATE: 0.0,
ModelParameter.LIST_SIZE: list_size,
ModelParameter.RANDOM_SEED: None
}
datasets = prepare_data(args, list_size, train_batch_size, eval_batch_size, True)
train_data, valid_data, test_data = datasets
if os.path.exists(args.checkpoint_dir):
log.info(f"Removing existing checkpoint directory: {args.checkpoint_dir}")
shutil.rmtree(args.checkpoint_dir, ignore_errors=True)
model = setup_model(model_params)
model(tf.data.experimental.get_single_element(train_data.take(1))[0], training=False)
model.summary(print_fn=log.info)
# Keras training
# model.fit(train_data, validation_data=valid_data, epochs=500)
# Custom loop
train_meta = {
'sample_pre_batch': False,
'max_epochs': tf.constant(args.max_epochs),
'step': tf.Variable(0, tf.int32),
'elapsed_time': tf.Variable(0., tf.float32),
'train_time': tf.Variable(0., tf.float32),
'valid_time': tf.Variable(0., tf.float32),
'secs_step': tf.Variable(0., tf.float32),
'train_loss': tf.keras.metrics.Mean(),
'metrics': {k: tf.keras.metrics.Mean() for k in (1, 5, 10)},
'best_result': tf.Variable(0., tf.float32)
}
ckpt = tf.train.Checkpoint(epoch=tf.Variable(0), optimizer=model.optimizer, model=model)
manager = tf.train.CheckpointManager(ckpt, args.checkpoint_dir, max_to_keep=1)
start_time = tf.timestamp()
for epoch in range(train_meta['max_epochs']):
train_meta['train_loss'].reset_states()
train(model, train_data, train_meta)
evaluate(model, valid_data, train_meta)
flag_best_result = ''
if train_meta['metrics'][5].result() > train_meta['best_result']:
manager.save()
train_meta['best_result'].assign(train_meta['metrics'][5].result())
flag_best_result = ' *'
train_meta['elapsed_time'].assign(tf.cast(tf.timestamp() - start_time, tf.float32))
ckpt.epoch.assign(epoch)
log.info(
f"epoch: {epoch+1:5d} "
f"step: {train_meta['step'].numpy():8d} "
f"elapsed time: {train_meta['elapsed_time'].numpy():7.2f}s "
f"train time: {train_meta['train_time'].numpy():6.2f}s "
f"secs/step: {train_meta['secs_step'].numpy():6.3f} "
f"val time: {train_meta['valid_time'].numpy():6.2f} "
f"train/loss: {train_meta['train_loss'].result():10.4f} "
f"val/ndcg@01: {train_meta['metrics'][1].result():10.4f} "
f"val/ndcg@05: {train_meta['metrics'][5].result():10.4f} "
f"val/ndcg@10: {train_meta['metrics'][10].result():10.4f}"
f"{flag_best_result}"
)
# Evaluate on the test data using the best model during training
ckpt = tf.train.Checkpoint(epoch=tf.Variable(0), model=model, optimizer=model.optimizer)
ckpt.restore(tf.train.latest_checkpoint(args.checkpoint_dir))
evaluate(model, test_data, train_meta)
log.info(
f"test/ndcg@01: {train_meta['metrics'][1].result():10.4f} "
f"test/ndcg@05: {train_meta['metrics'][5].result():10.4f} "
f"test/ndcg@10: {train_meta['metrics'][10].result():10.4f}"
)
# noinspection DuplicatedCode
[docs]def make_command_line_options():
cli = argparse.ArgumentParser()
cli.add_argument(
'--train-file',
required=True,
type=str,
help="The training tfrecords file."
)
cli.add_argument(
'--valid-file',
required=True,
type=str,
help="The validation tfrecords file."
)
cli.add_argument(
'--test-file',
required=True,
type=str,
help="The test tfrecords file."
)
cli.add_argument(
'--max-epochs',
required=True,
type=int,
help="The number of epochs to train for."
)
cli.add_argument(
'--checkpoint-dir',
required=True,
type=str,
help="The directory where model checkpoints will be saved."
)
cli.add_argument(
'--scaler',
required=False,
type=str,
nargs=2,
help=(
"This argument requires two parameters. The first is the path to "
"a scaler file created with the build dataset script. The second "
"is the name of the scaler to use. Choose one of: "
"minmax, standard, robust, power."
)
)
cli.add_argument(
'--random-seed',
required=False,
type=int,
help="The random seed to use for sampling query results."
)
cli.set_defaults(func=main)
return cli
if __name__ == '__main__':
clo = make_command_line_options()
cli_args = clo.parse_args()
cli_args.func(cli_args)