NVIDIA Deep Learning Examples
NVIDIA Deep Learning Examples
Wide & Deep for TensorFlow2
Resource
NVIDIA Deep Learning Examples
NVIDIA Deep Learning Examples
Wide & Deep for TensorFlow2

Wide & Deep Recommender model.

To train your model using the default parameters of the Wide & Deep model on the Outbrain dataset in TF32 or FP32 precision, perform the following steps. For the specifics concerning training and inference with custom settings, refer to the Advanced section.

  1. Clone the repository.
git clone https://github.com/NVIDIA/DeepLearningExamples
  1. Go to the WideAndDeep TensorFlow2 directory within the DeepLearningExamples repository:
cd DeepLearningExamples/TensorFlow2/Recommendation/WideAndDeep
  1. Download the Outbrain dataset.

The Outbrain dataset can be downloaded from Kaggle (requires a Kaggle account). Unzip the downloaded archive into the orig directory (for example, to /raid/outbrain/orig) and set the HOST_OUTBRAIN_PATH variable to the parent directory:

HOST_OUTBRAIN_PATH=/raid/outbrain
  1. Build the Wide & Deep Container.
cd DeepLearningExamples/TensorFlow2/Recommendation/WideAndDeep
docker build . -t wd2
  1. Preprocess the Outbrain dataset.

5.1. Start an interactive session in the Wide & Deep Container. Run preprocessing against the original Outbrain dataset to parquets. You can run preprocessing using NVTabular preprocessing (GPU).

docker run --runtime=nvidia --gpus=all --rm -it --ipc=host -v ${HOST_OUTBRAIN_PATH}:/outbrain wd2 bash

5.2. Start NVTabular GPU preprocessing. For more information, go to the Dataset preprocessing section.

bash scripts/preproc.sh

The result of preprocessing script is the NVTabular dataset stored in parquets. Files are generated in ${HOST_OUTBRAIN_PATH}/data.

  1. Train the model

6.1. Start an interactive session in the Wide & Deep Container

docker run --runtime=nvidia --gpus=all --rm -it --ipc=host -v ${HOST_OUTBRAIN_PATH}:/outbrain wd2 bash

6.2. Run training (${GPU} is an arbitrary number of gpu to be used)

horovodrun -np ${GPU} sh hvd_wrapper.sh python main.py

Training with Mixed Precision training with XLA:

horovodrun -np ${GPU} sh hvd_wrapper.sh python main.py --xla --amp

For complete usage, run:

python main.py -h
  1. Run validation or evaluation. If you want to run validation or evaluation, you can either:
  • use the checkpoint obtained from the training commands above, or
  • download the pretrained checkpoint from NGC.

In order to download the checkpoint from NGC, visit ngc.nvidia.com website and browse the available models. Download the checkpoint files and unzip them to some path, for example, to $HOST_OUTBRAIN_PATH/checkpoints/ (which is the default path for storing the checkpoints during training). The checkpoint requires around 700MB of disk space.

  1. Start validation/evaluation. In order to validate the checkpoint on the evaluation set, run the main.py script with the --evaluate and --use_checkpoint flags.
horovodrun -np ${GPU} sh hvd_wrapper.sh python main.py --evaluate --use_checkpoint

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark your performance to Training and evaluation performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.