NVIDIA Deep Learning Examples
NVIDIA Deep Learning Examples
VAE for TensorFlow1
Resource
NVIDIA Deep Learning Examples
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.

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

  2. Build the VAE TensorFlow NGC container.

    docker build . -t vae
    
  3. Launch the VAE-CF TensorFlow Docker container.

    docker run -it --rm --runtime=nvidia -v /data/vae-cf:/data vae /bin/bash
    
  4. 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>

  5. Prepare the dataset.

    python prepare_dataset.py
    
  6. 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
    
  7. 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