# 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.
# ================================
http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png

Exporting Ranking Models

In this example notebook we demonstrate how to export (save) NVTabular workflow and a ranking model for model deployment with Merlin Systems library.

Learning Objectives:

  • Export NVTabular workflow for model deployment

  • Export TensorFlow DLRM model for model deployment

We will follow the steps below:

  • Prepare the data with NVTabular and export NVTabular workflow

  • Train a DLRM model with Merlin Models and export the trained model

Importing Libraries

Let’s start with importing the libraries that we’ll use in this notebook.

import os
os.environ["TF_GPU_ALLOCATOR"]="cuda_malloc_async"

import nvtabular as nvt
from nvtabular.ops import *

from merlin.models.utils.example_utils import workflow_fit_transform

from merlin.schema.tags import Tags

import merlin.models.tf as mm
from merlin.io.dataset import Dataset
import tensorflow as tf

2022-04-05 16:15:14.772044: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-04-05 16:15:15.885699: I tensorflow/core/common_runtime/gpu/gpu_process_state.cc:214] Using CUDA malloc Async allocator for GPU: 0
2022-04-05 16:15:15.885845: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1525] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 16254 MB memory:  -> device: 0, name: Quadro GV100, pci bus id: 0000:15:00.0, compute capability: 7.0

Feature Engineering with NVTabular

We use the synthetic train and test datasets generated by mimicking the real Ali-CCP: Alibaba Click and Conversion Prediction dataset to build our recommender system ranking models.

If you would like to use real Ali-CCP dataset instead, you can download the training and test datasets on tianchi.aliyun.com. You can then use get_aliccp() function to curate the raw csv files and save them as the parquet.

from merlin.datasets.synthetic import generate_data

DATA_FOLDER = os.environ.get("DATA_FOLDER", "/workspace/data/")

train, valid = generate_data("aliccp-raw", 1000000, set_sizes=(0.7, 0.3))

# save the datasets as parquet files
train.to_ddf().to_parquet(os.path.join(DATA_FOLDER, 'train'))
valid.to_ddf().to_parquet(os.path.join(DATA_FOLDER, 'valid'))

(None,)

Let’s define our input and output paths.

train_path = os.path.join(DATA_FOLDER, "train", "part.0.parquet")
valid_path = os.path.join(DATA_FOLDER, "valid", "part.0.parquet")
output_path =os.path.join(DATA_FOLDER, "processed")

After we execute fit() and transform() functions on the raw dataset applying the operators defined in the NVTabular workflow pipeline below, the processed parquet files are saved to output_path.

%%time
user_id = ["user_id"] >> Categorify() >> TagAsUserID()
item_id = ["item_id"] >> Categorify() >> TagAsItemID()
targets = ["click"] >> AddMetadata(tags=[Tags.BINARY_CLASSIFICATION, "target"])

item_features = ["item_category", "item_shop", "item_brand"] >> Categorify() >> TagAsItemFeatures()

user_features = ['user_shops', 'user_profile', 'user_group', 
       'user_gender', 'user_age', 'user_consumption_2', 'user_is_occupied',
       'user_geography', 'user_intentions', 'user_brands', 'user_categories'] \
        >> Categorify() >> TagAsUserFeatures()

outputs = user_id + item_id + item_features + user_features + targets

workflow = nvt.Workflow(outputs)

train_dataset = nvt.Dataset(train_path)
valid_dataset = nvt.Dataset(valid_path)

workflow.fit(train_dataset)
workflow.transform(train_dataset).to_parquet(output_path=output_path + "/train/")
workflow.transform(valid_dataset).to_parquet(output_path=output_path + "/valid/")

/usr/local/lib/python3.8/dist-packages/cudf/core/dataframe.py:1253: UserWarning: The deep parameter is ignored and is only included for pandas compatibility.
  warnings.warn(

CPU times: user 1.25 s, sys: 736 ms, total: 1.99 s
Wall time: 2.05 s

We save NVTabular workflow model in the current working directory.

workflow.save('workflow')

Let’s check out our saved workflow model folder.

!apt-get install tree

!tree ./workflow

./workflow
├── categories
│   ├── unique.item_brand.parquet
│   ├── unique.item_category.parquet
│   ├── unique.item_id.parquet
│   ├── unique.item_shop.parquet
│   ├── unique.user_age.parquet
│   ├── unique.user_brands.parquet
│   ├── unique.user_categories.parquet
│   ├── unique.user_consumption_2.parquet
│   ├── unique.user_gender.parquet
│   ├── unique.user_geography.parquet
│   ├── unique.user_group.parquet
│   ├── unique.user_id.parquet
│   ├── unique.user_intentions.parquet
│   ├── unique.user_is_occupied.parquet
│   ├── unique.user_profile.parquet
│   └── unique.user_shops.parquet
├── metadata.json
└── workflow.pkl

1 directory, 18 files

Build and Train a DLRM model

In this example, we build, train, and export a Deep Learning Recommendation Model (DLRM) architecture. To learn more about how to train different deep learning models, how easily transition from one model to another and the seamless integration between data preparation and model training visit 03-Exploring-different-models.ipynb notebook.

NVTabular workflow above exports a schema file, schema.pbtxt, of our processed dataset. To learn more about the schema object, schema file and tags, you can explore 02-Merlin-Models-and-NVTabular-integration.ipynb.

# define train and valid dataset objects
train = Dataset(os.path.join(output_path, 'train', '*.parquet'))
valid = Dataset(os.path.join(output_path, 'valid', '*.parquet'))

# define schema object
schema = train.schema

target_column = schema.select_by_tag(Tags.TARGET).column_names[0]
target_column

'click'

model = mm.DLRMModel(
    schema,
    embedding_dim=64,
    bottom_block=mm.MLPBlock([128, 64]),
    top_block=mm.MLPBlock([128, 64, 32]),
    prediction_tasks=mm.BinaryClassificationTask(target_column, metrics=[tf.keras.metrics.AUC()])
)

%%time

model.compile('adam', run_eagerly=False)
model.fit(train, validation_data=valid, batch_size=16*1024)

2022-04-05 16:15:22.930133: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.

43/43 [==============================] - ETA: 0s - auc: 0.4981 - loss: 0.6932 - regularization_loss: 0.0000e+00 - total_loss: 0.6932

2022-04-05 16:15:31.611688: W tensorflow/core/grappler/optimizers/loop_optimizer.cc:907] Skipping loop optimization for Merge node with control input: cond/then/_0/cond/cond/branch_executed/_161

43/43 [==============================] - 8s 71ms/step - auc: 0.4981 - loss: 0.6932 - regularization_loss: 0.0000e+00 - total_loss: 0.6932 - val_auc: 0.5004 - val_loss: 0.6931 - val_regularization_loss: 0.0000e+00 - val_total_loss: 0.6931
CPU times: user 10.7 s, sys: 1.12 s, total: 11.8 s
Wall time: 9.41 s

<keras.callbacks.History at 0x7f6774260880>

Save model

The last step of machine learning (ML)/deep learning (DL) pipeline is to deploy the ETL workflow and saved model to production. In the production setting, we want to transform the input data as done during training (ETL). We need to apply the same mean/std for continuous features and use the same categorical mapping to convert the categories to continuous integer before we use the DL model for a prediction. Therefore, we deploy the NVTabular workflow with the Tensorflow model as an ensemble model to Triton Inference using Merlin Systems library very easily. The ensemble model guarantees that the same transformation is applied to the raw inputs.

Let’s save our DLRM model.

model.save('dlrm')

WARNING:absl:Found untraced functions such as sequential_block_8_layer_call_fn, sequential_block_8_layer_call_and_return_conditional_losses, binary_classification_task_layer_call_fn, binary_classification_task_layer_call_and_return_conditional_losses, click/binary_classification_task/output_layer_layer_call_fn while saving (showing 5 of 48). These functions will not be directly callable after loading.

INFO:tensorflow:Assets written to: dlrm/assets

INFO:tensorflow:Assets written to: dlrm/assets

We have NVTabular wokflow and DLRM model exported, now it is time to move on to the next step: model deployment with Merlin Systems.

Deploying the model with Merlin Systems

We trained and exported our ranking model and NVTabular workflow. In the next step, we will learn how to deploy our trained DLRM model into Triton Inference Server with Merlin Systems library. NVIDIA Triton Inference Server (TIS) simplifies the deployment of AI models at scale in production. TIS provides a cloud and edge inferencing solution optimized for both CPUs and GPUs. It supports a number of different machine learning frameworks such as TensorFlow and PyTorch.

For the next step, visit Merlin Systems library and execute Serving-Ranking-Models-With-Merlin-Systems notebook to deploy our saved DLRM and NVTabular workflow models as an ensemble to TIS and obtain prediction results for a qiven request. In doing so, you need to mount the saved DLRM and NVTabular workflow to the inference container following the instructions in the README.md.