The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

The main.py script provides an entry point to all the provided functionalities. This includes running training, testing and inference. The behavior of the script is controlled by command-line arguments listed below in the Parameters section. The prepare_dataset.py script can be used to preprocess the MovieLens 20m dataset.

Most of the deep learning logic is implemented in the vae/models subdirectory. The vae/load subdirectory contains the code for preprocessing the dataset. The vae/metrics subdirectory provides functions for computing the validation metrics such as recall and NDCG.

Parameters

The most important command-line parameters include:

• --data_dir which specifies the directory inside the docker container where the data will be stored, overriding the default location /data
• --checkpoint_dir which controls if and where the checkpoints will be stored
• --amp for enabling mixed precision training

There are also multiple parameters controlling the various hyperparameters of the training process, such as the learning rate, batch size etc.

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option, for example:

python main.py --help

usage: main.py [-h] [--train] [--test] [--inference_benchmark]
               [--amp] [--epochs EPOCHS]
               [--batch_size_train BATCH_SIZE_TRAIN]
               [--batch_size_validation BATCH_SIZE_VALIDATION]
               [--validation_step VALIDATION_STEP]
               [--warm_up_epochs WARM_UP_EPOCHS]
               [--total_anneal_steps TOTAL_ANNEAL_STEPS]
               [--anneal_cap ANNEAL_CAP] [--lam LAM] [--lr LR] [--beta1 BETA1]
               [--beta2 BETA2] [--top_results TOP_RESULTS] [--xla] [--trace]
               [--activation ACTIVATION] [--log_path LOG_PATH] [--seed SEED]
               [--data_dir DATA_DIR] [--checkpoint_dir CHECKPOINT_DIR]

Train a Variational Autoencoder for Collaborative Filtering in TensorFlow

optional arguments:
  -h, --help            show this help message and exit
  --train               Run training of VAE
  --test                Run validation of VAE
  --inference_benchmark
                        Benchmark the inference throughput and latency
  --amp          Enable Automatic Mixed Precision
  --epochs EPOCHS       Number of epochs to train
  --batch_size_train BATCH_SIZE_TRAIN
                        Global batch size for training
  --batch_size_validation BATCH_SIZE_VALIDATION
                        Used both for validation and testing
  --validation_step VALIDATION_STEP
                        Train epochs for one validation
  --warm_up_epochs WARM_UP_EPOCHS
                        Number of epochs to omit during benchmark
  --total_anneal_steps TOTAL_ANNEAL_STEPS
                        Number of annealing steps
  --anneal_cap ANNEAL_CAP
                        Annealing cap
  --lam LAM             Regularization parameter
  --lr LR               Learning rate
  --beta1 BETA1         Adam beta1
  --beta2 BETA2         Adam beta2
  --top_results TOP_RESULTS
                        Number of results to be recommended
  --xla                 Enable XLA
  --trace               Save profiling traces
  --activation ACTIVATION
                        Activation function
  --log_path LOG_PATH   Path to the detailed JSON log from to be created
  --seed SEED           Random seed for TensorFlow and numpy
  --data_dir DATA_DIR   Directory for storing the training data
  --checkpoint_dir CHECKPOINT_DIR
                        Path for saving a checkpoint after the training

Getting the data

The VA-CF model was trained on the MovieLens 20M dataset. The dataset can be preprocessed simply by running: python prepare_dataset.py in the Docker container. By default, the dataset will be stored in the /data directory. If you want to store the data in a different location, you can pass the desired location to the --data_dir argument.

Dataset guidelines

As a Collaborative Filtering model, VAE-CF only uses information about which user interacted with which item. For the MovieLens dataset, this means that a particular user has positively reviewed a particular movie. VAE-CF can be adapted to any other collaborative filtering task. The input to the model is generally a list of all interactions between users and items. One column of the CSV should contain user IDs, while the other should contain item IDs. Preprocessing for the MovieLens 20M dataset is provided in the vae/load/preprocessing.py file.

Training process

The training can be started by running the main.py script with the train argument. The resulting checkpoints containing the trained model weights are then stored in the directory specified by the --checkpoint_dir directory (by default no checkpoints are saved).

Additionally, a command-line argument called --results_dir (by default None) specifies where to save the following statistics in a JSON format:

1. a complete list of command-line arguments saved as <results_dir>/args.json, and
2. a dictionary of validation metrics and performance metrics recorded during training.

The main validation metric used is NDCG@100. Following the original VAE-CF paper we also report numbers for Recall@20 and Recall@50.

Multi-GPU training uses horovod.

Mixed precision support is controlled by the --amp command-line flag. It enables TensorFlow's Automatic Mixed Precision mode.

Inference process

Inference on a trained model can be run by passing the --inference_benchmark argument to the main.py script

python main.py --inference_benchmark [--amp] --checkpoint_dir ./checkpoints

This will generate a user with a collection of random items that they interacted with and run inference for that user multiple times to measure latency and throughput.