NVIDIA
NVIDIA
VAE for TensorFlow
Resource
NVIDIA
NVIDIA
VAE for TensorFlow

The Variational Autoencoder for collaborative filtering focuses on providing recommendations. This is an optimized implementation.

To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the VAE-CF model on the MovieLens 20m dataset. For the specifics concerning training and inference, see the Advanced section.

  1. Clone the repository. git clone https://github.com/NVIDIA/DeepLearningExamples cd DeepLearningExamples/Tensorflow/Recommendation/VAE_CF

2. Build the VAE TensorFlow NGC container.

```bash
docker build . -t vae
  1. Launch the VAE-CF TensorFlow Docker container.
docker run -it --rm --runtime=nvidia -v /data/vae-cf:/data vae /bin/bash
  1. Downloading the dataset: Here we use the MovieLens 20m dataset.
  • If you do not have the dataset downloaded: Run the commands below to download and extract the MovieLens dataset to the /data/ml-20m/extracted/ folder.
cd /data
mkdir ml-20m
cd ml-20m
mkdir extracted
cd extracted
wget http://files.grouplens.org/datasets/movielens/ml-20m.zip
unzip ml-20m.zip
  • If you already have the dataset downloaded and unzipped elsewhere: Run the below commands to first exit the current VAE-CF Docker container and then Restart the VAE-CF Docker Container (like in Step 3 above) by mounting the MovieLens dataset location
exit
docker run -it --rm --runtime=nvidia -v /data/vae-cf:/data -v :/data/ml-20m/extracted/ml-20m vae /bin/bash

where, the unzipped MovieLens dataset is at ``````

  1. Prepare the dataset.
python prepare_dataset.py
  1. Start training on 8 GPUs.
mpirun --bind-to numa --allow-run-as-root -np 8 -H localhost:8 python main.py --train --amp --checkpoint_dir ./checkpoints
  1. Start validation/evaluation.

The model is exported to the default model_dir and can be loaded and tested using:

python main.py --test --amp --checkpoint_dir ./checkpoints