# Copyright 2021 NVIDIA Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

Getting Started MovieLens: Training with PyTorch

Overview

We observed that PyTorch training pipelines can be slow as the dataloader is a bottleneck. The native dataloader in PyTorch randomly sample each item from the dataset, which is very slow. In our experiments, we are able to speed-up existing PyTorch pipelines using a highly optimized dataloader.

Applying deep learning models to recommendation systems faces unique challenges in comparison to other domains, such as computer vision and natural language processing. The datasets and common model architectures have unique characteristics, which require custom solutions. Recommendation system datasets have terabytes in size with billion examples but each example is represented by only a few bytes. For example, the Criteo CTR dataset, the largest publicly available dataset, is 1.3TB with 4 billion examples. The model architectures have normally large embedding tables for the users and items, which do not fit on a single GPU. You can read more in our blogpost.

Learning objectives

This notebook explains, how to use the NVTabular dataloader to accelerate PyTorch training.

1. Use NVTabular dataloader with PyTorch

2. Leverage multi-hot encoded input features

MovieLens25M

The MovieLens25M is a popular dataset for recommender systems and is used in academic publications. The dataset contains 25M movie ratings for 62,000 movies given by 162,000 users. Many projects use only the user/item/rating information of MovieLens, but the original dataset provides metadata for the movies, as well. For example, which genres a movie has. Although we may not improve state-of-the-art results with our neural network architecture, the purpose of this notebook is to explain how to integrate multi-hot categorical features into a neural network.

NVTabular dataloader for PyTorch

We’ve identified that the dataloader is one bottleneck in deep learning recommender systems when training pipelines with PyTorch. The dataloader cannot prepare the next batch fast enough, so and therefore, the GPU is not utilized.

As a result, we developed a highly customized tabular dataloader for accelerating existing pipelines in PyTorch. NVTabular dataloader’s features are:

• removing bottleneck of item-by-item dataloading

• enabling larger than memory dataset by streaming from disk

• reading data directly into GPU memory and remove CPU-GPU communication

• preparing batch asynchronously in GPU to avoid CPU-GPU communication

• supporting commonly used .parquet format for efficient data format

• easy integration into existing PyTorch pipelines by using similar API than the native one

# External dependencies
import os
import gc
import glob

import nvtabular as nvt

We define our base directory, containing the data.

INPUT_DATA_DIR = os.environ.get(
    "INPUT_DATA_DIR", os.path.expanduser("~/nvt-examples/movielens/data/")
)

Defining Hyperparameters

First, we define the data schema and differentiate between single-hot and multi-hot categorical features. Note, that we do not have any numerical input features.

BATCH_SIZE = 1024 * 32  # Batch Size
CATEGORICAL_COLUMNS = ["movieId", "userId"]  # Single-hot
CATEGORICAL_MH_COLUMNS = ["genres"]  # Multi-hot
NUMERIC_COLUMNS = []

# Output from ETL-with-NVTabular
TRAIN_PATHS = sorted(glob.glob(os.path.join(INPUT_DATA_DIR, "train", "*.parquet")))
VALID_PATHS = sorted(glob.glob(os.path.join(INPUT_DATA_DIR, "valid", "*.parquet")))

In the previous notebook, we used NVTabular for ETL and stored the workflow to disk. We can load the NVTabular workflow to extract important metadata for our training pipeline.

proc = nvt.Workflow.load(os.path.join(INPUT_DATA_DIR, "workflow"))

The embedding table shows the cardinality of each categorical variable along with its associated embedding size. Each entry is of the form (cardinality, embedding_size).

EMBEDDING_TABLE_SHAPES, MH_EMBEDDING_TABLE_SHAPES = nvt.ops.get_embedding_sizes(proc)
EMBEDDING_TABLE_SHAPES, MH_EMBEDDING_TABLE_SHAPES

{'genres': (21, 16), 'movieId': (56586, 512), 'userId': (162542, 512)}

Initializing NVTabular Dataloader for PyTorch

We import PyTorch and the NVTabular dataloader for PyTorch.

import torch
from nvtabular.loader.torch import TorchAsyncItr, DLDataLoader
from nvtabular.framework_utils.torch.models import Model
from nvtabular.framework_utils.torch.utils import process_epoch

First, we take a look on our dataloader and how the data is represented as tensors. The NVTabular dataloader are initialized as usually and we specify both single-hot and multi-hot categorical features as cats. The dataloader will automatically recognize the single/multi-hot columns and represent them accordingly.

# TensorItrDataset returns a single batch of x_cat, x_cont, y.

train_dataset = TorchAsyncItr(
    nvt.Dataset(TRAIN_PATHS),
    batch_size=BATCH_SIZE,
    cats=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,
    conts=NUMERIC_COLUMNS,
    labels=["rating"],
)
train_loader = DLDataLoader(
    train_dataset, batch_size=None, collate_fn=lambda x: x, pin_memory=False, num_workers=0
)

valid_dataset = TorchAsyncItr(
    nvt.Dataset(VALID_PATHS),
    batch_size=BATCH_SIZE,
    cats=CATEGORICAL_COLUMNS + CATEGORICAL_MH_COLUMNS,
    conts=NUMERIC_COLUMNS,
    labels=["rating"],
)
valid_loader = DLDataLoader(
    valid_dataset, batch_size=None, collate_fn=lambda x: x, pin_memory=False, num_workers=0
)

Let’s generate a batch and take a look on the input features. Normally, the NVTabular dataloader prepares batches of the structure (X_cat, X_cont, targets) - representing a Tensor for categorical input features, a Tensor for continous input features and the target. As we have both, single-hot categorical input features and multi-hot categorical input features, the first element of our batch is a tuple again.

The batch structure is ((X_cat_singlehot, X_cat_multihot), X_cont, targets). Let’s take a closer look on X_cat_multihot.

batch = next(iter(train_loader))
batch

((tensor([[ 11994, 124027],
          [  2550,  98809],
          [ 24262,  81377],
          ...,
          [    81,  66331],
          [  5292,  69908],
          [  1656, 100709]], device='cuda:0'),
  {'genres': (tensor([ 6,  9,  2,  ..., 18,  9, 16], device='cuda:0'),
    tensor([[    0],
            [    2],
            [    4],
            ...,
            [88432],
            [88435],
            [88437]], device='cuda:0'))}),
 None,
 tensor([1., 0., 1.,  ..., 1., 0., 1.], device='cuda:0'))

X_cat_multihot is a dictionary and each key contains a tuple of two Tensors. For the multi-hot categorical feature genres, we receive two Tensors values and nnzs.

X_cat_multihot = batch[0]
X_cat_multihot["genres"]

(tensor([ 6,  9,  2,  ..., 18,  9, 16], device='cuda:0'),
 tensor([[    0],
         [    2],
         [    4],
         ...,
         [88432],
         [88435],
         [88437]], device='cuda:0'))

values are the actual data, containing the genre IDs. Note that the Tensor has more values than the batch_size. The reason is, that one datapoint in the batch can contain more than one genre (multi-hot). nnzs are a supporting Tensor, describing the starting/ending point of associated values to each datapoint in the batch. For example, - if the first two values in nnzs is 0, 2, then the first 2 values (0, 1) in values are associated with the first datapoint in the batch (movieId/userId). - if the next value in nnzs is 6, then the 3rd, 4th and 5th value in values are associated with the second datapoint in the batch (continuing after the previous value stopped). - if the third value in nnzs is 7, then the 6th value in values are associated with the third datapoint in the batch. - and so on

X_cat_multihot["genres"][0].shape

torch.Size([88439])

X_cat_multihot["genres"][1].shape

torch.Size([32768, 1])

As each datapoint can have a different number of genres, it is more efficient to represent the genres as two flat tensors: One with the actual values (values) and one with the length for each datapoint (nnzs).

del batch
gc.collect()

80

Defining Neural Network Architecture

We implemented a simple PyTorch architecture.

• Single-hot categorical features are fed into an Embedding Layer

• Each value of a multi-hot categorical features is fed into an Embedding Layer and the multiple Embedding outputs are combined via summing

• The output of the Embedding Layers are concatenated

• The concatenated layers are fed through multiple feed-forward layers (Dense Layers, BatchNorm with ReLU activations)

You can see more details by checking out the implementation.

# ??Model

We initialize the model. EMBEDDING_TABLE_SHAPES needs to be a Tuple representing the cardinality for single-hot and multi-hot input features.

EMBEDDING_TABLE_SHAPES_TUPLE = (
    {
        CATEGORICAL_COLUMNS[0]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[0]],
        CATEGORICAL_COLUMNS[1]: EMBEDDING_TABLE_SHAPES[CATEGORICAL_COLUMNS[1]],
    },
    {CATEGORICAL_MH_COLUMNS[0]: MH_EMBEDDING_TABLE_SHAPES[CATEGORICAL_MH_COLUMNS[0]]},
)
EMBEDDING_TABLE_SHAPES_TUPLE

({'movieId': (56586, 512), 'userId': (162542, 512)}, {'genres': (21, 16)})

model = Model(
    embedding_table_shapes=EMBEDDING_TABLE_SHAPES_TUPLE,
    num_continuous=0,
    emb_dropout=0.0,
    layer_hidden_dims=[128, 128, 128],
    layer_dropout_rates=[0.0, 0.0, 0.0],
).to("cuda")
model

Model(
  (initial_cat_layer): ConcatenatedEmbeddings(
    (embedding_layers): ModuleList(
      (0): Embedding(56586, 512)
      (1): Embedding(162542, 512)
    )
    (dropout): Dropout(p=0.0, inplace=False)
  )
  (mh_cat_layer): MultiHotEmbeddings(
    (embedding_layers): ModuleList(
      (0): EmbeddingBag(21, 16, mode=sum)
    )
    (dropout): Dropout(p=0.0, inplace=False)
  )
  (initial_cont_layer): BatchNorm1d(0, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (layers): ModuleList(
    (0): Sequential(
      (0): Linear(in_features=1040, out_features=128, bias=True)
      (1): ReLU(inplace=True)
      (2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (3): Dropout(p=0.0, inplace=False)
    )
    (1): Sequential(
      (0): Linear(in_features=128, out_features=128, bias=True)
      (1): ReLU(inplace=True)
      (2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (3): Dropout(p=0.0, inplace=False)
    )
    (2): Sequential(
      (0): Linear(in_features=128, out_features=128, bias=True)
      (1): ReLU(inplace=True)
      (2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (3): Dropout(p=0.0, inplace=False)
    )
  )
  (output_layer): Linear(in_features=128, out_features=1, bias=True)
)

We initialize the optimizer.

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

We use the process_epoch function to train and validate our model. It iterates over the dataset and calculates as usually the loss and optimizer step.

%%time
from time import time

for epoch in range(1):
    start = time()
    train_loss, y_pred, y = process_epoch(train_loader,
                                          model,
                                          train=True,
                                          optimizer=optimizer)
    valid_loss, y_pred, y = process_epoch(valid_loader,
                                          model,
                                          train=False)
    print(f"Epoch {epoch:02d}. Train loss: {train_loss:.4f}. Valid loss: {valid_loss:.4f}.")
    t_final = time() - start
    total_rows = train_dataset.num_rows_processed + valid_dataset.num_rows_processed
    print(
        f"run_time: {t_final} - rows: {total_rows} - epochs: {epoch} - dl_thru: {total_rows / t_final}"
    )

Total batches: 611
Total batches: 152
Epoch 00. Train loss: 0.1916. Valid loss: 0.1703.
run_time: 30.9938702583313 - rows: 1530 - epochs: 0 - dl_thru: 49.36459974980791
CPU times: user 22.5 s, sys: 8.51 s, total: 31 s
Wall time: 31 s