VAE for TensorFlow1 | NVIDIA NGC

NVIDIA Deep Learning Examples

VAE for TensorFlow1

Resource

NVIDIA Deep Learning Examples

VAE for TensorFlow1

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.

Clone the repository. git clone https://github.com/NVIDIA/DeepLearningExamples cd DeepLearningExamples/Tensorflow/Recommendation/VAE_CF
Build the VAE TensorFlow NGC container.
```
docker build . -t vae
```

Launch the VAE-CF TensorFlow Docker container.

docker run -it --rm --runtime=nvidia -v /data/vae-cf:/data vae /bin/bash

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 <ml-20m folder path>:/data/ml-20m/extracted/ml-20m vae /bin/bash
```
  where, the unzipped MovieLens dataset is at <ml-20m folder path>
Prepare the dataset.
```
python prepare_dataset.py
```

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

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
```