Skip to main content
NVIDIA
NCF for PyTorch
Resource
NVIDIA
NCF for PyTorch

The NCF model focuses on providing recommendations. This is a modified implementation with improved overfitting and better accuracy.

Quick Start Guide

1. Build and launch an NCF PyTorch Docker container

After Docker is correctly set up, you can build the NCF image with:

docker build . -t nvidia_ncf

After that the NVIDIA NCF container can be launched with:

mkdir data
docker run --runtime=nvidia -it --rm --ipc=host  -v ${PWD}/data:/data nvidia_ncf bash

This will launch the container and mount the ./data directory as a volume to the /data directory inside the container. Any datasets and experiment results (logs, checkpoints etc.) saved to /data will be accessible in the './data' directory on the host.

2. Data preparation

Preprocessing consists of downloading the data, filtering out users that have less than 20 ratings (by default), sorting the data and dropping the duplicates. The preprocessed train and test data is then saved in PyTorch binary format to be loaded just before training.

No data augmentation techniques are used.

To download and preprocess the ml-20m dataset you can run:

./prepare_dataset.sh

Please note that this command will return immediately without downloading anything if the data is already present in the /data directory.

Other datasets

This implementation is tuned for the ml-20m and ml-1m datasets. Using other datasets might require tuning some hyperparameters (e.g., learning rate, beta1, beta2)

If you'd like to use your custom dataset you can do it by adding support for it in the prepare_dataset.sh and download_dataset.sh scripts. The required format of the data is a CSV file in which the first column contains the userID and the second column contains the itemID.

The performance of the model depends on the dataset size. Generally, the model should scale better for datasets containing more data points. For a smaller dataset the you might experience slower performance.

ml-1m

To download and preprocess the ml-1m dataset run:

./prepare_dataset.sh ml-1m

This will store the preprocessed training and evaluation data in the /data directory so that it can be later used to train the model (by passing the appropriate --data argument to the ncf.py script).

3. Run the training

After the docker container is launched, the training with the default hyperparameters can be started with:

./prepare_dataset.sh
python -m torch.distributed.launch --nproc_per_node=8 ncf.py --data /data/cache/ml-20m

This will result in a checkpoint file being written to /data/checkpoints/model.pth.

4. Test a trained model

The trained model can be evaluated by passing the --mode test flag to the run.sh script:

python -m torch.distributed.launch --nproc_per_node=8 ncf.py --data /data/cache/ml-20m  --mode test --checkpoint-path /data/checkpoints/model.pth

5. Hyperparameters and command line arguments

The default hyperparameters used are:

  • learning rate: 0.0045
  • beta1: 0.25
  • beta2: 0.5
  • training batch size: 1048576
  • epsilon: 1e-8
  • loss scale: 8192
  • negatives sampled for training: 4
  • use mixed precision training: Yes
  • number of GPUs used: 8

All these parameters can be controlled by passing command line arguments to the ncf.py script. To get a complete list of all command line arguments with descriptions and default values you can run:

python ncf.py --help
``

Modal Content

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.