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

NVIDIA uses cookies to improve your experience on our web site. We and our third-party partners also use cookies and other tools to collect and record information you provide as well as information about your interactions with our websites for performance improvement, analytics, and to assist in marketing efforts. By clicking "Accept All", you consent to our use of cookies and other tools as described in our Cookie Policy. You can manage your cookie settings by clicking on "Manage Settings." By continuing to use this site or by clicking one of the buttons below, you agree to our Terms of Service (which contains important waivers). Please see our Privacy Policy for more information on our privacy practices.