examples/scaling-criteo/03-Training-with-Merlin-Models-TensorFlow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a5b71acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1f95be1",
   "metadata": {},
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_merlin_scaling-criteo-03-training-with-merlin-models-tensorflow/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Scaling Criteo: Training with Merlin Models TensorFlow\n",
    "\n",
    "This notebook is created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container.\n",
    "\n",
    "\n",
    "## Overview\n",
    "\n",
    "The [Criteo 1TB Click Logs dataset](https://ailab.criteo.com/download-criteo-1tb-click-logs-dataset/) is a popular dataset in the recommender system community as it is one of the largest, public available dataset. It contains ~1.3 TB of uncompressed click logs containing over four billion samples spanning 24 days.\n",
    "\n",
    "We will train Facebook's [deep learning recommendation model (DLRM)](https://arxiv.org/abs/1906.00091) architecture with Merlin Models. We will assume you are familiar with Merlin Models' API and features. Otherwise, we recommend to start with the [Merlin Models examples](https://nvidia-merlin.github.io/models/stable/examples/index.html#).\n",
    "\n",
    "## Learning objectives\n",
    "\n",
    "- Train a DLRM architecture with Merlin Models on a large dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68941f39",
   "metadata": {},
   "source": [
    "## Training a DLRM model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db535904",
   "metadata": {},
   "source": [
    "Let's start with importing the libraries that we'll use in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3b347509",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3/dist-packages/requests/__init__.py:89: RequestsDependencyWarning: urllib3 (1.26.12) or chardet (3.0.4) doesn't match a supported version!\n",
      "  warnings.warn(\"urllib3 ({}) or chardet ({}) doesn't match a supported \"\n",
      "2022-11-29 17:20:01.942873: I tensorflow/core/platform/cpu_feature_guard.cc:194] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE3 SSE4.1 SSE4.2 AVX\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-11-29 17:20:06.579306: I tensorflow/core/common_runtime/gpu/gpu_process_state.cc:222] Using CUDA malloc Async allocator for GPU: 0\n",
      "2022-11-29 17:20:06.579506: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 16255 MB memory:  -> device: 0, name: Tesla V100-SXM2-32GB-LS, pci bus id: 0000:85:00.0, compute capability: 7.0\n",
      "2022-11-29 17:20:06.582211: I tensorflow/core/common_runtime/gpu/gpu_process_state.cc:222] Using CUDA malloc Async allocator for GPU: 1\n",
      "2022-11-29 17:20:06.582328: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 30655 MB memory:  -> device: 1, name: Tesla V100-SXM2-32GB-LS, pci bus id: 0000:8a:00.0, compute capability: 7.0\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"TF_GPU_ALLOCATOR\"] = \"cuda_malloc_async\"\n",
    "\n",
    "import glob\n",
    "import merlin.models.tf as mm\n",
    "from merlin.io.dataset import Dataset\n",
    "\n",
    "from merlin.schema import Tags\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b4923fa",
   "metadata": {},
   "source": [
    "Define the path to directories which contains the processed data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "71f448ee",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['/raid/data/criteo/test_dask/output/train/part_0.parquet',\n",
       "  '/raid/data/criteo/test_dask/output/train/part_1.parquet'],\n",
       " ['/raid/data/criteo/test_dask/output/valid/part_0.parquet',\n",
       "  '/raid/data/criteo/test_dask/output/valid/part_1.parquet'])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_path = os.environ.get(\"INPUT_DATA_DIR\", \"/raid/data/criteo/test_dask/output/\")\n",
    "\n",
    "# path to processed data\n",
    "PATH_TO_TRAIN_DATA = sorted(glob.glob(os.path.join(input_path, \"train\", \"*.parquet\")))\n",
    "PATH_TO_VALID_DATA = sorted(glob.glob(os.path.join(input_path, \"valid\", \"*.parquet\")))\n",
    "\n",
    "PATH_TO_TRAIN_DATA, PATH_TO_VALID_DATA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74aa21e3",
   "metadata": {},
   "source": [
    "We define some hyperparameters for the model architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "54c35685",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = int(os.environ.get(\"BATCH_SIZE\", 64 * 1024))\n",
    "EMBEDDING_SIZE = 8\n",
    "EPOCHS = 1\n",
    "LR = 0.01\n",
    "OPTIMIZER = tf.keras.optimizers.SGD(learning_rate=LR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97375902",
   "metadata": {},
   "source": [
    "We will use Merlin Dataset object to initialize the dataloaders. It provides a dataset schema to initialize the model architectures. The [Merlin Models examples](https://nvidia-merlin.github.io/models/stable/examples/index.html#) will explain more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1aab6cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = Dataset(PATH_TO_TRAIN_DATA, part_mem_fraction=0.08)\n",
    "valid = Dataset(PATH_TO_VALID_DATA, part_mem_fraction=0.08)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edf40eb9",
   "metadata": {},
   "source": [
    "We initialize the DLRM architecture with Merlin Models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ce4117ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = mm.DLRMModel(\n",
    "    train.schema,\n",
    "    embedding_dim=EMBEDDING_SIZE,\n",
    "    bottom_block=mm.MLPBlock([128, EMBEDDING_SIZE]),\n",
    "    top_block=mm.MLPBlock([128, 64, 32]),\n",
    "    prediction_tasks=mm.BinaryClassificationTask(\n",
    "        train.schema.select_by_tag(Tags.TARGET).column_names[0]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "598a56da",
   "metadata": {},
   "source": [
    "We compile and train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "52ea88cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2867/2867 [==============================] - 224s 72ms/step - loss: 0.1563 - precision: 0.0378 - recall: 2.0745e-04 - binary_accuracy: 0.9677 - auc: 0.5673 - regularization_loss: 0.0000e+00 - loss_batch: 0.1563 - val_loss: 0.1385 - val_precision: 0.0000e+00 - val_recall: 0.0000e+00 - val_binary_accuracy: 0.9680 - val_auc: 0.6279 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 0.1405\n",
      "CPU times: user 8min 32s, sys: 13min 13s, total: 21min 45s\n",
      "Wall time: 3min 54s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f58457b1c40>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "model.compile(optimizer=OPTIMIZER, run_eagerly=False)\n",
    "model.fit(train,\n",
    "          validation_data=valid,\n",
    "          batch_size=BATCH_SIZE,\n",
    "          epochs=EPOCHS\n",
    "          )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04f3cf26",
   "metadata": {},
   "source": [
    "## Evaluate the model\n",
    "\n",
    "Finally, we can evaluate our model on the validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "db25ed20",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2787/2787 [==============================] - 89s 31ms/step - loss: 0.1385 - precision: 0.0000e+00 - recall: 0.0000e+00 - binary_accuracy: 0.9680 - auc: 0.6279 - regularization_loss: 0.0000e+00 - loss_batch: 0.1385\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': 0.13853003084659576,\n",
       " 'precision': 0.0,\n",
       " 'recall': 0.0,\n",
       " 'binary_accuracy': 0.9679937958717346,\n",
       " 'auc': 0.6279259324073792,\n",
       " 'regularization_loss': 0.0,\n",
       " 'loss_batch': 0.14048266410827637}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_metrics = model.evaluate(valid, batch_size=BATCH_SIZE, return_dict=True)\n",
    "eval_metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79c3d245",
   "metadata": {},
   "source": [
    "## Save the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b58b134",
   "metadata": {},
   "source": [
    "We save the model to disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b0f9e68c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Unsupported signature for serialization: ((PredictionOutput(predictions=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/predictions'), targets=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/targets'), positive_item_ids=None, label_relevant_counts=None, valid_negatives_mask=None, negative_item_ids=None, sample_weight=None), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f57c5f8ecd0>), {}).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Function `_wrapped_model` contains input name(s) C1, C10, C11, C12, C13, C14, C15, C16, C17, C18, C19, C2, C20, C21, C22, C23, C24, C25, C26, C3, C4, C5, C6, C7, C8, C9, I1, I10, I11, I12, I13, I2, I3, I4, I5, I6, I7, I8, I9 with unsupported characters which will be renamed to c1, c10, c11, c12, c13, c14, c15, c16, c17, c18, c19, c2, c20, c21, c22, c23, c24, c25, c26, c3, c4, c5, c6, c7, c8, c9, i1, i10, i11, i12, i13, i2, i3, i4, i5, i6, i7, i8, i9 in the SavedModel.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Unsupported signature for serialization: ((PredictionOutput(predictions=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/predictions'), targets=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/targets'), positive_item_ids=None, label_relevant_counts=None, valid_negatives_mask=None, negative_item_ids=None, sample_weight=None), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f57c5f8ecd0>), {}).\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Unsupported signature for serialization: ((PredictionOutput(predictions=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/predictions'), targets=TensorSpec(shape=(None, 1), dtype=tf.float32, name='outputs/targets'), positive_item_ids=None, label_relevant_counts=None, valid_negatives_mask=None, negative_item_ids=None, sample_weight=None), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7f57c5f8ecd0>), {}).\n",
      "WARNING:absl:Found untraced functions such as train_compute_metrics, model_context_layer_call_fn, model_context_layer_call_and_return_conditional_losses, output_layer_layer_call_fn, output_layer_layer_call_and_return_conditional_losses while saving (showing 5 of 161). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /raid/data/criteo/test_dask/output/dlrm/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: /raid/data/criteo/test_dask/output/dlrm/assets\n"
     ]
    }
   ],
   "source": [
    "model.save(os.path.join(input_path, \"dlrm\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3bdb1e2",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "We trained Facebook's popular DLRM architecture with only ~5 commands on the large criteo dataset.   \n",
    "\n",
    "## Next steps\n",
    "\n",
    "The next step  is to [deploy the NVTabular workflow and DLRM model](04-Triton-Inference-with-Merlin-Models-TensorFlow.ipynb) to production.\n",
    "\n",
    "If you are interested more in different architecture and training models with Merlin Models, we recommend to check out our [Merlin Models examples](https://nvidia-merlin.github.io/models/stable/examples/index.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}