run.py

#!/usr/bin/python3
import click
import os
from functools import partial
import vae.config
from vae.load.generic import load_dataset
from vae.utils.round import round_8
from vae.metrics.recall import recall
from vae.metrics.ndcg import ndcg


@click.command()
@click.option('--train', is_flag=True, help='Run training of VAE')
@click.option('--test', is_flag=True, help='Run validation of VAE')
@click.option('--benchmark', is_flag=True, help='Run benchmark of VAE')
@click.option('--use_tf_amp', is_flag=True,
              help='Enable Automatic Mixed Precision to speedup fp32 computation using tensor cores')
@click.option('--dataset', default=vae.config.ML_20M, help='Dataset to use')
@click.option('--gpu_number', default=0, help='Number of GPU used in training or validation')
@click.option('--number_of_gpus', default=1, help='How many GPUs to use during training or validation')
@click.option('--number_of_epochs', default=200, help='Number of epochs to train')
@click.option('--batch_size_train', default=10000)
@click.option('--batch_size_validation', default=10000, help='Used both for validation and testing')
@click.option('--validation_step', default=5)
@click.option('--warm_up_epochs', default=5, help='Number of epochs to omit during benchmark')
@click.option('--total_anneal_steps', default=200000, help='Number of annealing steps')
@click.option('--anneal_cap', default=0.2, help='Annealing cap')
@click.option('--lam', default=1e-2, help='Regularization parameter')
@click.option('--lr', default=1e-3, help='Learning rate')
def main(train,
         test,
         benchmark,
         dataset,
         gpu_number,
         number_of_gpus,
         number_of_epochs,
         batch_size_train,
         batch_size_validation,
         use_tf_amp,
         validation_step,
         warm_up_epochs,
         total_anneal_steps,
         anneal_cap,
         lam,
         lr):
    if not train and not test and not benchmark:
        print("Choose one or more:")
        for option in "train", "test", "benchmark":
            print("\t" + option)
        exit(1)

    if train and benchmark:
        print("Use of train and benchmark together is not allowed")
        exit(1)

    # Suppress TF warnings
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    # set AMP
    os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1' if use_tf_amp else '0'

    # load dataset
    (train_data,
     validation_data_input,
     validation_data_true,
     test_data_input,
     test_data_true) = load_dataset(dataset,
                                    val_ratio=0.0731651997,
                                    test_ratio=0.0731651997,
                                    min_items=5,
                                    min_users=0,
                                    seed=0)

    # make sure all dims and sizes are divisible by 8
    number_of_train_users, number_of_items = train_data.shape
    number_of_items = round_8(number_of_items)

    for data in [train_data,
                 validation_data_input,
                 validation_data_true,
                 test_data_input,
                 test_data_true]:
        number_of_users, _ = data.shape
        data.resize(number_of_users, number_of_items)

    # compute VAE dims
    number_of_users, number_of_items = train_data.shape
    encoder_dims = [number_of_items, 600, 200]

    # create VAE
    if number_of_gpus == 0:
        from vae.models.Mult_VAE_training import VAE
        vae = VAE(train_data,
                  encoder_dims,
                  total_anneal_steps=total_anneal_steps,
                  anneal_cap=anneal_cap,
                  batch_size_train=batch_size_train,
                  batch_size_validation=batch_size_validation,
                  lam=lam,
                  lr=lr,
                  device='/device:CPU')
    elif number_of_gpus == 1:
        from vae.models.Mult_VAE_training import VAE
        vae = VAE(train_data,
                  encoder_dims,
                  total_anneal_steps=total_anneal_steps,
                  anneal_cap=anneal_cap,
                  batch_size_train=batch_size_train,
                  batch_size_validation=batch_size_validation,
                  lam=lam,
                  lr=lr,
                  device='/device:GPU:' + str(gpu_number))
    else:
        batches_per_epoch = -(-number_of_train_users // batch_size_train)  # ceil div
        if (batches_per_epoch % number_of_gpus != 0):
            print("Number of batches must be divisible by number of GPUs")
            exit(1)
        from vae.models.Mult_VAE_training_horovod import VAE_Horovod
        vae = VAE_Horovod(train_data,
                          encoder_dims,
                          total_anneal_steps=total_anneal_steps,
                          anneal_cap=anneal_cap,
                          batch_size_train=batch_size_train,
                          batch_size_validation=batch_size_validation,
                          lam=lam,
                          lr=lr)

    metrics = {'ndcg@100': partial(ndcg, R=100),
               'recall@20': partial(recall, R=20),
               'recall@50': partial(recall, R=50)}

    if train:
        vae.train(n_epochs=number_of_epochs,
                  train_data=train_data,
                  validation_data_input=test_data_input,
                  validation_data_true=test_data_true,
                  batch_size_train=batch_size_train,
                  batch_size_validation=batch_size_validation,
                  metrics=metrics,
                  validation_step=validation_step)
        if number_of_gpus > 1:
            print("Saving is not supported with horovod multigpu yet")
        else:
            vae.save()

    if benchmark:
        vae.benchmark(n_epochs=number_of_epochs,
                      warm_up_epochs=warm_up_epochs,
                      train_data=train_data,
                      validation_data_input=test_data_input,
                      validation_data_true=test_data_true,
                      batch_size_train=batch_size_train,
                      batch_size_validation=batch_size_validation,
                      metrics=metrics,
                      validation_step=validation_step)
        vae.save()

    if test and number_of_gpus <= 1:
        vae.load()

        test_results = \
            vae.test(test_data_input=test_data_input,
                     test_data_true=test_data_true,
                     metrics=metrics)

        for k, v in test_results.items():
            print("{}:\t{}".format(k, v))
    elif test and number_of_gpus > 1:
        print("Testing is not supported with horovod multigpu yet")


if __name__ == '__main__':
    main()