notebooks/embedding_plugin.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HugeCTR Embedding  Plugin for TensorFlow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook introduces a TensorFlow (TF) plugin for the HugeCTR embedding layer, embedding_plugin, where users may benefit from both the computational efficiency of the HugeCTR embedding layer and the ease of use of TensorFlow (TF)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build embedding_plugin ##\n",
    "**1.**<br>\n",
    "Before you can use the embedding_plugin, you must first build HugeCTR. You can do so by running the following commands:\n",
    "```shell\n",
    "$ git clone https://github.com/NVIDIA/HugeCTR.git\n",
    "$ cd HugeCTR\n",
    "$ git submodule update --init --recursive\n",
    "$ mkdir -p build && cd build\n",
    "$ cmake -DCMAKE_BUILD_TYPE=Release -DSM=80 .. # target is NVIDIA A100\n",
    "$ make -j$(nproc)\n",
    "```\n",
    "A dynamic library is generated in the `lib/` directory that you'll have to load using TensorFlow. You can directly import `hugectr_tf_ops.py`, where we prepare the codes to load that dynamic library and wrap some operations for convenient usage, in your python script to be used with the embedding_plugin."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Verify Accuracy ##\n",
    "To verify whether the embedding_plugin can obtain the correct result, you can generate synthetic data for testing purposes as shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run this cell to clear all variables.\n",
    "%reset -f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import tensorflow and some modules\n",
    "import tensorflow as tf\n",
    "# do not let TF allocate all GPU memory\n",
    "devices = tf.config.list_physical_devices(\"GPU\")\n",
    "for dev in devices:\n",
    "    tf.config.experimental.set_memory_growth(dev, True)\n",
    "    \n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import hugectr_tf_ops.py to use embedding_plugin ops\n",
    "import sys\n",
    "sys.path.append(\"../tools/embedding_plugin/python/\")\n",
    "import hugectr_tf_ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init embedding table value:\n",
      " [[ 1.  2.  3.  4.]\n",
      " [ 5.  6.  7.  8.]\n",
      " [ 9. 10. 11. 12.]\n",
      " [13. 14. 15. 16.]\n",
      " [17. 18. 19. 20.]\n",
      " [21. 22. 23. 24.]\n",
      " [25. 26. 27. 28.]\n",
      " [29. 30. 31. 32.]]\n"
     ]
    }
   ],
   "source": [
    "# generate a random embedding table and show\n",
    "vocabulary_size = 8\n",
    "slot_num = 3\n",
    "embedding_vector_size = 4\n",
    "\n",
    "table = np.float32([i for i in range(1, vocabulary_size * embedding_vector_size + 1)]).reshape(vocabulary_size, embedding_vector_size)\n",
    "print(\"init embedding table value:\\n\", table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In HugeCTR, the corresponding dense shape of the input keys is `[batch_size, slot_num, max_nnz]`, and `0` is a valid key. Therefore, `-1` is used to denote invalid keys, which only occupy that position in the corresponding dense keys matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "the dense shape of inputs keys: (4, 3, 2)\n"
     ]
    }
   ],
   "source": [
    "# generate random keys to lookup from embedding table.\n",
    "keys = np.array([[[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [2,  6]],  # nnz = 2\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [-1, -1]], # nnz = 0\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [6, -1]],  # nnz = 1\n",
    "                 \n",
    "                 [[0, -1],   # nnz = 1\n",
    "                  [1, -1],   # nnz = 1\n",
    "                  [2, -1]]], # nnz = 1\n",
    "                dtype=np.int64) \n",
    "print(\"the dense shape of inputs keys:\", keys.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [34. 36. 38. 40.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 0.  0.  0.  0.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [25. 26. 27. 28.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 9. 10. 11. 12.]]], shape=(4, 3, 4), dtype=float32)\n",
      "\n",
      "\n",
      "second forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [33.800995   35.800995   37.800995   39.800995  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 0.          0.          0.          0.        ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [24.900497   25.900497   26.900497   27.900497  ]]\n",
      "\n",
      " [[ 0.90024936  1.9002494   2.9002495   3.9002495 ]\n",
      "  [ 4.9002495   5.9002495   6.9002495   7.9002495 ]\n",
      "  [ 8.900497    9.900497   10.900497   11.900497  ]]], shape=(4, 3, 4), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# define a simple forward propagation and backward propagation with embedding_plugin\n",
    "# NOTE: cause hugectr_tf_ops.init can only be called once, if you want to run this cell multi-times, please restart the kernel.\n",
    "\n",
    "with tf.GradientTape() as tape:\n",
    "    # hugectr_tf_ops embedding_plugin initialize\n",
    "    hugectr_tf_ops.init(visiable_gpus=[0], seed=123, key_type='int64', value_type='float', batch_size=4, batch_size_eval=4)\n",
    "    \n",
    "    # create a embedding_layer with embedding_plugin\n",
    "    embedding_name = hugectr_tf_ops.create_embedding(init_value=table, opt_hparams=[0.1, 0.9, 0.99, 1e-3], \n",
    "                                              name_='embedding_verification', \n",
    "                                              max_vocabulary_size_per_gpu=vocabulary_size,\n",
    "                                              slot_num=slot_num, embedding_vec_size=embedding_vector_size,\n",
    "                                              embedding_type='distributed', max_nnz=2)\n",
    "    \n",
    "    # convert dense input keys to SparseTensor\n",
    "    indices = tf.where(keys != -1)\n",
    "    values = tf.gather_nd(keys, indices)\n",
    "    \n",
    "    # create a Variable used in backward propagation\n",
    "    bp_trigger = tf.Variable(initial_value=1.0, trainable=True, dtype=tf.float32)\n",
    "    \n",
    "    # get forward result\n",
    "    forward_result = hugectr_tf_ops.fprop(embedding_name=embedding_name,\n",
    "                                   sparse_indices=indices, values=values, dense_shape=keys.shape,\n",
    "                                   output_type=tf.float32, is_training=True, bp_trigger=bp_trigger)\n",
    "    print(\"forward_result:\\n\", forward_result)\n",
    "    \n",
    "    # compute gradients & update params\n",
    "    grads = tape.gradient(forward_result, bp_trigger)\n",
    "    \n",
    "    # do second forward propagation to check whether embedding table is updated.\n",
    "    forward_2 = hugectr_tf_ops.fprop(embedding_name=embedding_name,\n",
    "                              sparse_indices=indices, values=values, dense_shape=keys.shape,\n",
    "                              output_type=tf.float32, is_training=True, bp_trigger=bp_trigger)\n",
    "    print(\"\\n\")\n",
    "    print(\"second forward_result:\\n\", forward_2)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [34. 36. 38. 40.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 0.  0.  0.  0.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [25. 26. 27. 28.]]\n",
      "\n",
      " [[ 1.  2.  3.  4.]\n",
      "  [ 5.  6.  7.  8.]\n",
      "  [ 9. 10. 11. 12.]]], shape=(4, 3, 4), dtype=float32)\n",
      "\n",
      "\n",
      "tf second forward_result:\n",
      " tf.Tensor(\n",
      "[[[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [33.800995   35.800995   37.800995   39.800995  ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [ 0.          0.          0.          0.        ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [24.900497   25.900497   26.900497   27.900497  ]]\n",
      "\n",
      " [[ 0.90024906  1.9002491   2.900249    3.900249  ]\n",
      "  [ 4.900249    5.900249    6.900249    7.900249  ]\n",
      "  [ 8.900497    9.900497   10.900497   11.900497  ]]], shape=(4, 3, 4), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "# similarly, use original tensorflow op to compare whether results are consistent.\n",
    "\n",
    "# define a tf embedding layer\n",
    "class EmbeddingLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, vocabulary_size, embedding_vec_size,\n",
    "                init_value):\n",
    "        super(EmbeddingLayer, self).__init__()\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.init_value = init_value\n",
    "        \n",
    "    def build(self, _):\n",
    "        self.Var = self.add_weight(shape=(self.vocabulary_size, self.embedding_vec_size),\n",
    "                                         initializer=tf.constant_initializer(value=self.init_value))\n",
    "        \n",
    "    def call(self, inputs):\n",
    "        return tf.nn.embedding_lookup_sparse(self.Var, inputs, sp_weights=None, combiner=\"sum\")\n",
    "    \n",
    "with tf.GradientTape() as tape:\n",
    "    # reshape keys into [batch_size * slot_num, max_nnz]\n",
    "    reshape_keys = np.reshape(keys, newshape=(-1, keys.shape[-1]))\n",
    "    indices = tf.where(reshape_keys != -1)\n",
    "    values = tf.gather_nd(reshape_keys, indices)\n",
    "\n",
    "    # define a layer\n",
    "    tf_layer = EmbeddingLayer(vocabulary_size, embedding_vector_size, table)\n",
    "    \n",
    "    # wrap input keys components into a SparseTensor\n",
    "    sparse_tensor = tf.sparse.SparseTensor(indices, values, reshape_keys.shape)\n",
    "    \n",
    "    tf_forward = tf_layer(sparse_tensor)\n",
    "    print(\"tf forward_result:\\n\", tf.reshape(tf_forward, [keys.shape[0], keys.shape[1], tf_forward.shape[-1]]))\n",
    "    \n",
    "    # define an optimizer\n",
    "    optimizer = tf.keras.optimizers.Adam(learning_rate=0.1, beta_1=0.9, beta_2=0.99, epsilon=1e-3)\n",
    "    \n",
    "    # compute gradients & update params\n",
    "    grads = tape.gradient(tf_forward, tf_layer.trainable_weights)\n",
    "    optimizer.apply_gradients(zip(grads, tf_layer.trainable_weights))\n",
    "    \n",
    "    # do second forward propagation to check whether params are updated.\n",
    "    tf_forward_2 = tf_layer(sparse_tensor)\n",
    "    print(\"\\n\")\n",
    "    print(\"tf second forward_result:\\n\", tf.reshape(tf_forward_2, [keys.shape[0], keys.shape[1], tf_forward_2.shape[-1]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Consistent in first forward propagation? True\n",
      "Consistent in second forward propagation? True\n"
     ]
    }
   ],
   "source": [
    "# assert whether embedding_plugin's results are consistent with tensorflow original ops\n",
    "first_forward_consistent = np.allclose(forward_result.numpy(), \n",
    "                                tf.reshape(tf_forward, [keys.shape[0], keys.shape[1], tf_forward.shape[-1]]).numpy())\n",
    "print(\"Consistent in first forward propagation?\", first_forward_consistent)\n",
    "\n",
    "second_forwad_consistent = np.allclose(forward_2.numpy(), \n",
    "                                tf.reshape(tf_forward_2, [keys.shape[0], keys.shape[1], tf_forward_2.shape[-1]]))\n",
    "print(\"Consistent in second forward propagation?\", second_forwad_consistent)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results from embedding_plugins and original TF ops are consistent in both first and second forward propagation, which means the embedding_plugin can get the same forward result and perform the same backward propagation as TF ops. Therefore, the embedding_plugin can obtain the correct results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DeepFM demo ##\n",
    "In this notebook, TF 2.x is used to build the DeepFM model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Models with the Embedding_Plugin ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To proceed, Kernel must be restarted.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first, import tensorflow and import plugin ops from hugectr_tf_ops.py\n",
    "import tensorflow as tf\n",
    "# do not let TF allocate all GPU memory\n",
    "devices = tf.config.list_physical_devices(\"GPU\")\n",
    "for dev in devices:\n",
    "    tf.config.experimental.set_memory_growth(dev, True)\n",
    "import sys\n",
    "sys.path.append(\"../tools/embedding_plugin/python/\")\n",
    "import hugectr_tf_ops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wrap plugin ops into a TF layer for easy use\n",
    "class PluginEmbedding(tf.keras.layers.Layer):\n",
    "    def __init__(self,\n",
    "                 vocabulary_size,\n",
    "                 slot_num,\n",
    "                 embedding_vec_size,\n",
    "                 gpu_count,\n",
    "                 initializer=False,\n",
    "                 name='plugin_embedding',\n",
    "                 embedding_type='localized',\n",
    "                 optimizer='Adam',\n",
    "                 opt_hparam=[0.1, 0.9, 0.99, 1e-3],\n",
    "                 update_type='Local',\n",
    "                 atomic_update=True,\n",
    "                 max_feature_num=int(1e3),\n",
    "                 max_nnz=1,\n",
    "                 combiner='sum',\n",
    "                 ):\n",
    "        super(PluginEmbedding, self).__init__()\n",
    "\n",
    "        self.vocabulary_size_each_gpu = (vocabulary_size // gpu_count) + 1 \n",
    "        self.slot_num = slot_num\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.embedding_type = embedding_type\n",
    "        self.optimizer_type = optimizer\n",
    "        self.opt_hparam = opt_hparam\n",
    "        self.update_type = update_type\n",
    "        self.atomic_update = atomic_update\n",
    "        self.max_feature_num = max_feature_num\n",
    "        self.max_nnz = max_nnz\n",
    "        self.combiner = combiner\n",
    "        self.gpu_count = gpu_count\n",
    "\n",
    "        self.name_ = hugectr_tf_ops.create_embedding(initializer, name_=name, embedding_type=self.embedding_type, \n",
    "                                             optimizer_type=self.optimizer_type, \n",
    "                                             max_vocabulary_size_per_gpu=self.vocabulary_size_each_gpu,\n",
    "                                             opt_hparams=self.opt_hparam, update_type=self.update_type,\n",
    "                                             atomic_update=self.atomic_update, slot_num=self.slot_num,\n",
    "                                             max_nnz=self.max_nnz, max_feature_num=self.max_feature_num,\n",
    "                                             embedding_vec_size=self.embedding_vec_size, \n",
    "                                             combiner=self.combiner)\n",
    "\n",
    "    def build(self, _):\n",
    "        self.bp_trigger = self.add_weight(name=\"bp_trigger\",\n",
    "                                          shape=(1,), dtype=tf.float32, trainable=True)\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, row_offsets, value_tensors, nnz_array, output_shape, training=False):\n",
    "        return hugectr_tf_ops.fprop_v3(embedding_name=self.name_, row_offsets=row_offsets, value_tensors=value_tensors, \n",
    "                                nnz_array=nnz_array, bp_trigger=self.bp_trigger, is_training=training,\n",
    "                                output_shape=output_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define other TF layers\n",
    "class Multiply(tf.keras.layers.Layer):\n",
    "    def __init__(self, out_units):\n",
    "        super(Multiply, self).__init__()\n",
    "        self.out_units = out_units\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self.w = self.add_weight(name='weight_vector', shape=(input_shape[1], self.out_units),\n",
    "                                 initializer='glorot_uniform', trainable=True)\n",
    "    \n",
    "    def call(self, inputs):\n",
    "        return inputs * self.w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build DeepFM with plugin layer\n",
    "class DeepFM_PluginEmbedding(tf.keras.models.Model):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size, \n",
    "                 embedding_vec_size,\n",
    "                 which_embedding,\n",
    "                 dropout_rate, # list of float\n",
    "                 deep_layers, # list of int\n",
    "                 initializer,\n",
    "                 gpus,\n",
    "                 batch_size,\n",
    "                 batch_size_eval,\n",
    "                 embedding_type = 'localized',\n",
    "                 slot_num=1,\n",
    "                 seed=123):\n",
    "        super(DeepFM_PluginEmbedding, self).__init__()\n",
    "        tf.keras.backend.clear_session()\n",
    "        tf.compat.v1.set_random_seed(seed)\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.which_embedding = which_embedding\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.deep_layers = deep_layers\n",
    "        self.gpus = gpus\n",
    "        self.batch_size = batch_size\n",
    "        self.batch_size_eval = batch_size_eval \n",
    "        self.slot_num = slot_num\n",
    "        self.embedding_type = embedding_type\n",
    "\n",
    "        if isinstance(initializer, str):\n",
    "            initializer = False\n",
    "            \n",
    "        # when building model with embedding_plugin ops, init() should be called prior to any other ops.\n",
    "        hugectr_tf_ops.init(visiable_gpus=gpus, seed=seed, key_type='int64', value_type='float', \n",
    "                        batch_size=batch_size, batch_size_eval=batch_size_eval)\n",
    "        \n",
    "        # create a embedding_plugin layer\n",
    "        self.plugin_embedding_layer = PluginEmbedding(vocabulary_size=vocabulary_size, slot_num=slot_num, \n",
    "                                            embedding_vec_size=embedding_vec_size + 1, \n",
    "                                            embedding_type=embedding_type,\n",
    "                                            gpu_count=len(gpus), initializer=initializer)\n",
    "        \n",
    "        # other layers with TF original ops\n",
    "        self.deep_dense = []\n",
    "        for i, deep_units in enumerate(self.deep_layers):\n",
    "            self.deep_dense.append(tf.keras.layers.Dense(units=deep_units, activation=None, use_bias=True,\n",
    "                                                         kernel_initializer='glorot_normal', \n",
    "                                                         bias_initializer='glorot_normal'))\n",
    "            self.deep_dense.append(tf.keras.layers.Dropout(dropout_rate[i]))\n",
    "        self.deep_dense.append(tf.keras.layers.Dense(units=1, activation=None, use_bias=True,\n",
    "                                                     kernel_initializer='glorot_normal',\n",
    "                                                     bias_initializer=tf.constant_initializer(0.01)))\n",
    "        self.add_layer = tf.keras.layers.Add()\n",
    "        self.y_act = tf.keras.layers.Activation(activation='sigmoid')\n",
    "\n",
    "        self.dense_multi = Multiply(1)\n",
    "        self.dense_embedding = Multiply(self.embedding_vec_size)\n",
    "\n",
    "        self.concat_1 = tf.keras.layers.Concatenate()\n",
    "        self.concat_2 = tf.keras.layers.Concatenate()\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, dense_feature, sparse_feature, training=True):\n",
    "        \"\"\"\n",
    "        forward propagation.\n",
    "        #arguments:\n",
    "            dense_feature: [batch_size, dense_dim]\n",
    "            sparse_feature: for OriginalEmbedding, it is a SparseTensor, and the dense shape is [batch_size * slot_num, max_nnz];\n",
    "                            for PluginEmbedding, it is a list of [row_offsets, value_tensors, nnz_array]. \n",
    "        \"\"\"\n",
    "        with tf.name_scope(\"embedding_and_slice\"):\n",
    "            dense_0 = tf.cast(tf.expand_dims(dense_feature, 2), dtype=tf.float32) # [batchsize, dense_dim, 1]\n",
    "            dense_mul = self.dense_multi(dense_0) # [batchsize, dense_dim, 1]\n",
    "            dense_emb = self.dense_embedding(dense_0) # [batchsize, dense_dim, embedding_vec_size]\n",
    "            dense_mul = tf.reshape(dense_mul, [dense_mul.shape[0], -1]) # [batchsize, dense_dim * 1]\n",
    "            dense_emb = tf.reshape(dense_emb, [dense_emb.shape[0], -1]) # [batchsize, dense_dim * embedding_vec_size]\n",
    "\n",
    "            sparse = self.plugin_embedding_layer(sparse_feature[0], sparse_feature[1], sparse_feature[2],\n",
    "                                                output_shape=[self.batch_size, self.slot_num, self.embedding_vec_size + 1],\n",
    "                                                training=training) # [batch_size, self.slot_num, self.embedding_vec_size + 1]\n",
    "\n",
    "            sparse_1 = tf.slice(sparse, [0, 0, self.embedding_vec_size], [-1, self.slot_num, 1]) #[batchsize, slot_num, 1]\n",
    "            sparse_1 = tf.squeeze(sparse_1, 2) # [batchsize, slot_num]\n",
    "\n",
    "            sparse_emb = tf.slice(sparse, [0, 0, 0], [-1, self.slot_num, self.embedding_vec_size]) #[batchsize, slot_num, embedding_vec_size]\n",
    "            sparse_emb = tf.reshape(sparse_emb, [-1, self.slot_num * self.embedding_vec_size]) #[batchsize, slot_num * embedding_vec_size]\n",
    "        \n",
    "        with tf.name_scope(\"FM\"):\n",
    "            with tf.name_scope(\"first_order\"):\n",
    "                first = self.concat_1([dense_mul, sparse_1]) # [batchsize, dense_dim + slot_num]\n",
    "                first_out = tf.reduce_sum(first, axis=-1, keepdims=True) # [batchsize, 1]\n",
    "                \n",
    "            with tf.name_scope(\"second_order\"):\n",
    "                hidden = self.concat_2([dense_emb, sparse_emb]) # [batchsize, (dense_dim + slot_num) * embedding_vec_size]\n",
    "                second = tf.reshape(hidden, [-1, dense_feature.shape[1] + self.slot_num, self.embedding_vec_size])\n",
    "                square_sum = tf.math.square(tf.math.reduce_sum(second, axis=1, keepdims=True)) # [batchsize, 1, embedding_vec_size]\n",
    "                sum_square = tf.math.reduce_sum(tf.math.square(second), axis=1, keepdims=True) # [batchsize, 1, embedding_vec_size]\n",
    "                \n",
    "                second_out = 0.5 * (sum_square - square_sum) # [batchsize, 1, embedding_vec_size]\n",
    "                second_out = tf.math.reduce_sum(second_out, axis=-1, keepdims=False) # [batchsize, 1]\n",
    "                \n",
    "        with tf.name_scope(\"Deep\"):\n",
    "            for i, layer in enumerate(self.deep_dense):\n",
    "                if i % 2 == 0: # dense\n",
    "                    hidden = layer(hidden)\n",
    "                else: # dropout\n",
    "                    hidden = layer(hidden, training)\n",
    "\n",
    "        y = self.add_layer([hidden, first_out, second_out])\n",
    "        y = self.y_act(y) # [batchsize, 1]\n",
    "\n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above cells wrap the embedding_plugin ops into a TF layer, and uses that layer to define a TF DeepFM model. Similarly, define an embedding layer with TF original ops, and define a DeepFM model with that layer. Because embedding_plugin supports model parallelism, the parameters of the original TF embedding layer are equally distributed to each GPU for a fair performance comparison."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Models with the Original TF Ops ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a TF embedding layer with TF original ops\n",
    "class OriginalEmbedding(tf.keras.layers.Layer):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size,\n",
    "                 embedding_vec_size,\n",
    "                 initializer='uniform',\n",
    "                 combiner=\"sum\",\n",
    "                 gpus=[0]):\n",
    "        super(OriginalEmbedding, self).__init__()\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size \n",
    "        if isinstance(initializer, str):\n",
    "            self.initializer = tf.keras.initializers.get(initializer)\n",
    "        else:\n",
    "            self.initializer = initializer\n",
    "        if combiner not in [\"sum\", \"mean\"]:\n",
    "            raise RuntimeError(\"combiner must be one of \\{'sum', 'mean'\\}.\")\n",
    "        self.combiner = combiner\n",
    "        if (not isinstance(gpus, list)) and (not isinstance(gpus, tuple)):\n",
    "            raise RuntimeError(\"gpus must be a list or tuple.\")\n",
    "        self.gpus = gpus\n",
    "\n",
    "    def build(self, _):\n",
    "        if isinstance(self.initializer, tf.keras.initializers.Initializer):\n",
    "            if len(self.gpus) > 1:\n",
    "                self.embeddings_params = list()\n",
    "                mod_size = self.vocabulary_size % len(self.gpus)\n",
    "                vocabulary_size_each_gpu = [(self.vocabulary_size // len(self.gpus)) + (1 if dev_id < mod_size else 0)\n",
    "                                            for dev_id in range(len(self.gpus))]\n",
    "\n",
    "                for i, gpu in enumerate(self.gpus):\n",
    "                    with tf.device(\"/gpu:%d\" %gpu):\n",
    "                        params_i = self.add_weight(name=\"embedding_\" + str(gpu), \n",
    "                                                   shape=(vocabulary_size_each_gpu[i], self.embedding_vec_size),\n",
    "                                                   initializer=self.initializer)\n",
    "                    self.embeddings_params.append(params_i)\n",
    "\n",
    "            else:\n",
    "                self.embeddings_params = self.add_weight(name='embeddings', \n",
    "                                                        shape=(self.vocabulary_size, self.embedding_vec_size),\n",
    "                                                        initializer=self.initializer)\n",
    "        else:\n",
    "            self.embeddings_params = self.initializer\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, keys, output_shape):\n",
    "        result = tf.nn.embedding_lookup_sparse(self.embeddings_params, keys, \n",
    "                                             sp_weights=None, combiner=self.combiner)\n",
    "        return tf.reshape(result, output_shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define DeepFM model with original TF embedding layer\n",
    "class DeepFM_OriginalEmbedding(tf.keras.models.Model):\n",
    "    def __init__(self, \n",
    "                 vocabulary_size, \n",
    "                 embedding_vec_size,\n",
    "                 which_embedding,\n",
    "                 dropout_rate, # list of float\n",
    "                 deep_layers, # list of int\n",
    "                 initializer,\n",
    "                 gpus,\n",
    "                 batch_size,\n",
    "                 batch_size_eval,\n",
    "                 embedding_type = 'localized',\n",
    "                 slot_num=1,\n",
    "                 seed=123):\n",
    "        super(DeepFM_OriginalEmbedding, self).__init__()\n",
    "        tf.keras.backend.clear_session()\n",
    "        tf.compat.v1.set_random_seed(seed)\n",
    "\n",
    "        self.vocabulary_size = vocabulary_size\n",
    "        self.embedding_vec_size = embedding_vec_size\n",
    "        self.which_embedding = which_embedding\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.deep_layers = deep_layers\n",
    "        self.gpus = gpus\n",
    "        self.batch_size = batch_size\n",
    "        self.batch_size_eval = batch_size_eval \n",
    "        self.slot_num = slot_num\n",
    "        self.embedding_type = embedding_type\n",
    "\n",
    "        self.original_embedding_layer = OriginalEmbedding(vocabulary_size=vocabulary_size, \n",
    "                                            embedding_vec_size=embedding_vec_size + 1, \n",
    "                                            initializer=initializer, gpus=gpus)\n",
    "        self.deep_dense = []\n",
    "        for i, deep_units in enumerate(self.deep_layers):\n",
    "            self.deep_dense.append(tf.keras.layers.Dense(units=deep_units, activation=None, use_bias=True,\n",
    "                                                         kernel_initializer='glorot_normal', \n",
    "                                                         bias_initializer='glorot_normal'))\n",
    "            self.deep_dense.append(tf.keras.layers.Dropout(dropout_rate[i]))\n",
    "        self.deep_dense.append(tf.keras.layers.Dense(units=1, activation=None, use_bias=True,\n",
    "                                                     kernel_initializer='glorot_normal',\n",
    "                                                     bias_initializer=tf.constant_initializer(0.01)))\n",
    "        self.add_layer = tf.keras.layers.Add()\n",
    "        self.y_act = tf.keras.layers.Activation(activation='sigmoid')\n",
    "\n",
    "        self.dense_multi = Multiply(1)\n",
    "        self.dense_embedding = Multiply(self.embedding_vec_size)\n",
    "\n",
    "        self.concat_1 = tf.keras.layers.Concatenate()\n",
    "        self.concat_2 = tf.keras.layers.Concatenate()\n",
    "\n",
    "    @tf.function\n",
    "    def call(self, dense_feature, sparse_feature, training=True):\n",
    "        \"\"\"\n",
    "        forward propagation.\n",
    "        #arguments:\n",
    "            dense_feature: [batch_size, dense_dim]\n",
    "            sparse_feature: for OriginalEmbedding, it is a SparseTensor, and the dense shape is [batch_size * slot_num, max_nnz];\n",
    "                            for PluginEmbedding, it is a list of [row_offsets, value_tensors, nnz_array]. \n",
    "        \"\"\"\n",
    "        with tf.name_scope(\"embedding_and_slice\"):\n",
    "            dense_0 = tf.cast(tf.expand_dims(dense_feature, 2), dtype=tf.float32) # [batchsize, dense_dim, 1]\n",
    "            dense_mul = self.dense_multi(dense_0) # [batchsize, dense_dim, 1]\n",
    "            dense_emb = self.dense_embedding(dense_0) # [batchsize, dense_dim, embedding_vec_size]\n",
    "            dense_mul = tf.reshape(dense_mul, [dense_mul.shape[0], -1]) # [batchsize, dense_dim * 1]\n",
    "            dense_emb = tf.reshape(dense_emb, [dense_emb.shape[0], -1]) # [batchsize, dense_dim * embedding_vec_size]\n",
    "\n",
    "            sparse = self.original_embedding_layer(sparse_feature, output_shape=[-1, self.slot_num, self.embedding_vec_size + 1])\n",
    "\n",
    "            sparse_1 = tf.slice(sparse, [0, 0, self.embedding_vec_size], [-1, self.slot_num, 1]) #[batchsize, slot_num, 1]\n",
    "            sparse_1 = tf.squeeze(sparse_1, 2) # [batchsize, slot_num]\n",
    "\n",
    "            sparse_emb = tf.slice(sparse, [0, 0, 0], [-1, self.slot_num, self.embedding_vec_size]) #[batchsize, slot_num, embedding_vec_size]\n",
    "            sparse_emb = tf.reshape(sparse_emb, [-1, self.slot_num * self.embedding_vec_size]) #[batchsize, slot_num * embedding_vec_size]\n",
    "        \n",
    "        with tf.name_scope(\"FM\"):\n",
    "            with tf.name_scope(\"first_order\"):\n",
    "                first = self.concat_1([dense_mul, sparse_1]) # [batchsize, dense_dim + slot_num]\n",
    "                first_out = tf.reduce_sum(first, axis=-1, keepdims=True) # [batchsize, 1]\n",
    "                \n",
    "            with tf.name_scope(\"second_order\"):\n",
    "                hidden = self.concat_2([dense_emb, sparse_emb]) # [batchsize, (dense_dim + slot_num) * embedding_vec_size]\n",
    "                second = tf.reshape(hidden, [-1, dense_feature.shape[1] + self.slot_num, self.embedding_vec_size])\n",
    "                square_sum = tf.math.square(tf.math.reduce_sum(second, axis=1, keepdims=True)) # [batchsize, 1, embedding_vec_size]\n",
    "                sum_square = tf.math.reduce_sum(tf.math.square(second), axis=1, keepdims=True) # [batchsize, 1, embedding_vec_size]\n",
    "                \n",
    "                second_out = 0.5 * (sum_square - square_sum) # [batchsize, 1, embedding_vec_size]\n",
    "                second_out = tf.math.reduce_sum(second_out, axis=-1, keepdims=False) # [batchsize, 1]\n",
    "                \n",
    "        with tf.name_scope(\"Deep\"):\n",
    "            for i, layer in enumerate(self.deep_dense):\n",
    "                if i % 2 == 0: # dense\n",
    "                    hidden = layer(hidden)\n",
    "                else: # dropout\n",
    "                    hidden = layer(hidden, training)\n",
    "\n",
    "        y = self.add_layer([hidden, first_out, second_out])\n",
    "        y = self.y_act(y) # [batchsize, 1]\n",
    "\n",
    "        return y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dataset is needed to use these models for training. [Kaggle Criteo datasets](http://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset/) provided by CriteoLabs is used as the training dataset. The original training set contains 45,840,617 examples. Each example contains a label (0 by default or 1 if the ad was clicked) and 39 features in which 13 of them are integer and the other 26 are categorial. Since TFRecord is suitable for the training process and the Criteo dataset is missing numerous values across the feature columns, preprocessing is needed. The original test set won't be used because it doesn't contain labels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dataset processing ###\n",
    "1. Download dataset from [http://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset/](http://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset/).\n",
    "2. Extract the dataset by running the following command. \n",
    "    ```shell\n",
    "    $ tar zxvf dac.tar.gz\n",
    "    ```\n",
    "3. Preprocess the datast and set missing values.\n",
    "Preprocessing functions are defined in [preprocess.py](../tools/embedding_plugin/performance_profile/preprocess.py). Open that file and check the codes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify source csv name and output csv name, run this command will do the preprocessing.\n",
    "# Warning: this command will take serveral hours to do preprocessing.\n",
    "%run ../tools/embedding_plugin/performance_profile/preprocess.py \\\n",
    "    --src_csv_path=train.txt --dst_csv_path=train.out.txt \\\n",
    "    --normalize_dense=0 --feature_cross=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Split the dataset by running the following commands:\n",
    "```shell\n",
    "$ head -n 36672493 train.out.txt > train\n",
    "$ tail -n 9168124 train.out.txt > valtest\n",
    "$ head -n 4584062 valtest > val\n",
    "$ tail -n 4584062 valtest > test\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. Convert the dataset into a TFRecord file. Converting functions are defined in [txt2tfrecord.py](../tools/embedding_plugin/performance_profile/txt2tfrecord.py). Open that file and check the codes.\n",
    "After the data preprocessing is completed, *.tfrecord file(s) will be generated, which can be used for training. The training loop can now be configured to use the dataset and models to perform the training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# specify source name and output tfrecord name, run this command will do the converting.\n",
    "# Warning: this command will take half an hour to do converting.\n",
    "%run ../tools/embedding_plugin/performance_profile/txt2tfrecord.py \\\n",
    "    --src_txt_name=train --dst_tfrecord_name=train.tfrecord \\\n",
    "    --normalized=0 --use_multi_process=1 --shard_num=1 \n",
    "    # if multi tfrecord files are wanted, set shard_num to the number of files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define training loop and do training ###\n",
    "In [read_data.py](../tools/embedding_plugin/performance_profile/read_data.py), some preprocessing and TF data reading pipeline creation functions are defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set env path, so that some modules can be imported\n",
    "sys.path.append(\"../tools/embedding_plugin/performance_profile/\")\n",
    "\n",
    "import txt2tfrecord as utils\n",
    "from read_data import create_dataset\n",
    "import time\n",
    "import logging\n",
    "logging.basicConfig(format='%(asctime)s %(message)s')\n",
    "logging.root.setLevel('INFO')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose wich model for training\n",
    "which_model = \"Plugin\" # change it to \"Original\", if you want to try the model define with original tf ops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set some hyper parameters for training process\n",
    "if (\"Plugin\" == which_model):\n",
    "    batch_size = 16384\n",
    "    n_epochs = 1\n",
    "    distribute_keys = 1 \n",
    "    gpus = [0] # use GPU0\n",
    "    embedding_type = 'distributed'\n",
    "    vocabulary_size = 1737710\n",
    "    embedding_vec_size = 10\n",
    "    slot_num = 26\n",
    "    batch_size_eval = 1 * len(gpus)\n",
    "    \n",
    "elif (\"Original\" == which_model):\n",
    "    batch_size = 16384\n",
    "    n_epochs = 1\n",
    "    distribute_keys = 0\n",
    "    gpus = [0] # use GPU0\n",
    "    vocabulary_size = 1737710\n",
    "    embedding_vec_size = 10\n",
    "    slot_num = 26\n",
    "    batch_size_eval = 1 * len(gpus)\n",
    "    embedding_type = 'distributed'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define feature_description to read tfrecord examples.\n",
    "cols = [utils.idx2key(idx, False) for idx in range(0, utils.NUM_TOTAL_COLUMNS)]\n",
    "feature_desc = dict()\n",
    "for col in cols:\n",
    "    if col == 'label' or col.startswith(\"I\"):\n",
    "        feature_desc[col] = tf.io.FixedLenFeature([], tf.int64) # scaler\n",
    "    else: \n",
    "        feature_desc[col] = tf.io.FixedLenFeature([1], tf.int64) # [slot_num, nnz]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# please set data_path to your tfrecord\n",
    "data_path = \"../tools/embedding_plugin/performance_profile/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/parallel_for/pfor.py:2380: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-11-23 09:45:31,393 From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/parallel_for/pfor.py:2380: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.\n"
     ]
    }
   ],
   "source": [
    "# create tfrecord reading pipeling\n",
    "dataset_names = [data_path + \"./train.tfrecord\"]\n",
    "dataset = create_dataset(dataset_names=dataset_names,\n",
    "                         feature_desc=feature_desc,\n",
    "                         batch_size=batch_size,\n",
    "                         n_epochs=n_epochs,\n",
    "                         distribute_keys=tf.constant(distribute_keys != 0, dtype=tf.bool),\n",
    "                         gpu_count=len(gpus),\n",
    "                         embedding_type=tf.constant(embedding_type, dtype=tf.string))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define loss function and optimizer used in other TF layers.\n",
    "optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)\n",
    "loss_fn = tf.keras.losses.BinaryCrossentropy(from_logits=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model instance\n",
    "if \"Original\" == which_model:\n",
    "    model = DeepFM_OriginalEmbedding(vocabulary_size=vocabulary_size, embedding_vec_size=embedding_vec_size, \n",
    "                       which_embedding=which_model, embedding_type=embedding_type,\n",
    "                       dropout_rate=[0.5] * 10, deep_layers=[1024] * 10,\n",
    "                       initializer='uniform', gpus=gpus, batch_size=batch_size, batch_size_eval=batch_size_eval,\n",
    "                       slot_num=slot_num)\n",
    "elif \"Plugin\" == which_model:\n",
    "    model = DeepFM_PluginEmbedding(vocabulary_size=vocabulary_size, embedding_vec_size=embedding_vec_size, \n",
    "                       which_embedding=which_model, embedding_type=embedding_type,\n",
    "                       dropout_rate=[0.5] * 10, deep_layers=[1024] * 10,\n",
    "                       initializer='uniform', gpus=gpus, batch_size=batch_size, batch_size_eval=batch_size_eval,\n",
    "                       slot_num=slot_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define training step\n",
    "@tf.function\n",
    "def _train_step(dense_batch, sparse_batch, y_batch, model, loss_fn, optimizer):\n",
    "    with tf.GradientTape() as tape:\n",
    "        y_batch = tf.cast(y_batch, dtype=tf.float32)\n",
    "        logits = model(dense_batch, sparse_batch, training=True)\n",
    "        loss = loss_fn(y_batch, logits)\n",
    "        loss /= dense_batch.shape[0]\n",
    "    grads = tape.gradient(loss, model.trainable_weights)\n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_weights))\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2020-11-23 09:47:33,030 begin to train\n",
      "2020-11-23 09:47:45,262 step: 100, loss: 0.0002282, elapsed time: 12.23093 seconds.\n",
      "2020-11-23 09:47:54,933 step: 200, loss: 0.0002397, elapsed time: 9.67155 seconds.\n",
      "2020-11-23 09:48:04,670 step: 300, loss: 0.0002279, elapsed time: 9.73684 seconds.\n",
      "2020-11-23 09:48:14,446 step: 400, loss: 0.0002361, elapsed time: 9.77609 seconds.\n",
      "2020-11-23 09:48:24,239 step: 500, loss: 0.0002272, elapsed time: 9.79255 seconds.\n",
      "2020-11-23 09:48:34,095 step: 600, loss: 0.0002441, elapsed time: 9.85616 seconds.\n",
      "2020-11-23 09:48:43,973 step: 700, loss: 0.0002372, elapsed time: 9.87768 seconds.\n",
      "2020-11-23 09:48:53,886 step: 800, loss: 0.0002601, elapsed time: 9.91286 seconds.\n",
      "2020-11-23 09:49:03,821 step: 900, loss: 0.0002376, elapsed time: 9.93518 seconds.\n",
      "2020-11-23 09:49:13,772 step: 1000, loss: 0.0002412, elapsed time: 9.95164 seconds.\n",
      "2020-11-23 09:49:23,746 step: 1100, loss: 0.0002340, elapsed time: 9.97380 seconds.\n",
      "2020-11-23 09:49:33,740 step: 1200, loss: 0.0002556, elapsed time: 9.99350 seconds.\n",
      "2020-11-23 09:49:43,747 step: 1300, loss: 0.0002328, elapsed time: 10.00727 seconds.\n",
      "2020-11-23 09:49:53,767 step: 1400, loss: 0.0002415, elapsed time: 10.02032 seconds.\n",
      "2020-11-23 09:50:03,804 step: 1500, loss: 0.0002307, elapsed time: 10.03668 seconds.\n",
      "2020-11-23 09:50:13,876 step: 1600, loss: 0.0002254, elapsed time: 10.07223 seconds.\n",
      "2020-11-23 09:50:23,967 step: 1700, loss: 0.0002475, elapsed time: 10.09131 seconds.\n",
      "2020-11-23 09:50:34,061 step: 1800, loss: 0.0002527, elapsed time: 10.09329 seconds.\n",
      "2020-11-23 09:50:44,161 step: 1900, loss: 0.0002480, elapsed time: 10.09980 seconds.\n",
      "2020-11-23 09:50:54,262 step: 2000, loss: 0.0002428, elapsed time: 10.10191 seconds.\n",
      "2020-11-23 09:51:04,371 step: 2100, loss: 0.0002515, elapsed time: 10.10869 seconds.\n",
      "2020-11-23 09:51:14,474 step: 2200, loss: 0.0002502, elapsed time: 10.10321 seconds.\n",
      "2020-11-23 09:51:18,241 Train End. Elapsed Time: 225.210 seconds.\n"
     ]
    }
   ],
   "source": [
    "# training loop\n",
    "logging.info(\"begin to train\")\n",
    "begin_time = time.time()\n",
    "display_begin = begin_time\n",
    "for step, datas in enumerate(dataset):\n",
    "    label, dense, others = datas[0], datas[1], datas[2:]\n",
    "    if tf.constant(distribute_keys != 0, dtype=tf.bool):\n",
    "        sparse = others[0:3]\n",
    "    else:\n",
    "        sparse = others[-1]\n",
    "    \n",
    "    train_loss = _train_step(dense, sparse, label, model, loss_fn, optimizer)\n",
    "    loss_value = train_loss.numpy()\n",
    "    \n",
    "    if (step % 100 == 0 and step != 0):\n",
    "        display_end = time.time()\n",
    "        logging.info(\"step: %d, loss: %.7f, elapsed time: %.5f seconds.\" %(step, loss_value, (display_end - display_begin)))\n",
    "        display_begin = display_end\n",
    "        \n",
    "end_time = time.time()\n",
    "logging.info(\"Train End. Elapsed Time: %.3f seconds.\" %(end_time - begin_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## API signature ##\n",
    "All embedding_plugin APIs are defined in [hugectr_tf_ops.py](../tools/embedding_plugin/python/hugectr_tf_ops.py)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>\n",
       "table {float:left}\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "<style>\n",
    "table {float:left}\n",
    "</style>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  init(visiable_gpus, seed=0, key_type='int64', value_type='float', batch_size=1, batch_size_eval=1)\n",
    "  ```\n",
    "  \n",
    "This function is used to create resource manager, which manages resources used by embedding_plugin.\n",
    "**IMPORTANT:** This function can only be called once. It must be called before any other embedding_plugin API is called Currently, only key_type='int64', value_type='float' has been tested.\n",
    "\n",
    "\n",
    "| Args ||\n",
    "| :-----| :---- |\n",
    "| visiable_gpus | list of integers, used to specify which gpus will be used by embedding_plugin. |\n",
    "| seed | integer, the initializer random seed for embedding_plugin. |\n",
    "| key_type| string, can be one of {'uint32', 'int64'}. Used to specify the input keys data type. |\n",
    "| value_type| string, can be one of {'float', 'half'}. Used to specify the data type of embedding_plugin forward result. |\n",
    "| batch_size| integer, batch_size used in training process. |\n",
    "| batch_size_eval| integer, batch_size used in evaluation process. |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  embedding_name = create_embedding(init_value, name_='hugectr_embedding', embedding_type='localized',\n",
    "                                     optimizer_type='Adam', max_vocabulary_size_per_gpu=1, slot_size_array=[],\n",
    "                                     opt_hparams=[0.001], update_type='Local', atomic_update=true, scaler=1.0,\n",
    "                                     slot_num=1, max_nnz=1, max_feature_num=1000000, embedding_vec_size=1,\n",
    "                                     combiner='sum')\n",
    "  ```\n",
    "  \n",
    "| Args ||\n",
    "| :-----| :---- |\n",
    "|init_value| can be a `bool` or a 2-D matrix with `dtype=tf.float32`. When it is `bool`, parameters will be randomly initialized. When it is a 2-D matrix with `dtype=tf.float32`, that matrix will be used to initialize parameters, and the matrix's row-index will be deemed to be key of the embedding table.|\n",
    "|name_|string, the name of this embedding layer. If `name_` is unique, then it will be used as the embedding layer name, otherwise, numerical suffix will be automatically added to `name_` to form an unique name for this embedding layer. |\n",
    "|embedding_type| string, can be one of {'localized', 'distributed'}. |\n",
    "| optimizer_type| string, can be one of {'Adam', 'MomentumSGD', 'Nesterov', 'SGD'}. | \n",
    "|max_vocabulary_size_per_gpu| integer, used to allocate GPU memory spaces for embedding layer.|\n",
    "|slot_size_array| list of integers, used to allocate GPU memory spaces precisely for embedding layer.|\n",
    "|opt_hparams| list of floats, used to specify hyper parameters for optimizer.<br>For `Adam`, `opt_hparams` must be a list of `[learning_rate, beta1, beta2, epsilon]`.<br>For `MomentumSGD`, `opt_hparams` must be a list of `[learning_rate, momentum_factor]`.<br>For `Nesterov`, `opt_hparams` must be a list of `[learning_rate, momentum_factor]`.<br>For `SGD`, `opt_hparams` must be a list of `[learning_rate]`.|\n",
    "|update_type| string, can be one of {'Local', 'Global', 'LazyGlobal'}. |\n",
    "|atomic_update| bool, only used in `SGD` optimizer. |\n",
    "|scaler| float, can be one of {1.0, 128.0, 256.0, 512.0, 1024.0}, used in `mixed_precission` training. |\n",
    "|slot_num| integer, how many slots (feature-fields) are unified in a single embedding layer. |\n",
    "|max_nnz| integer, the number of valid keys in a single slot.|\n",
    "|max_feature_num| integer, the number of valid keys in a single input sample.|\n",
    "|embedding_vec_size| integer, the embedding vector size of this embedding layer.|\n",
    "|combier|string, can be one of {'mean', 'sum'}. specify how to combine different embedding vector in the same slot.|\n",
    "\n",
    "|Returns||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor, dtype=tf.string. An unique name for this embedding layer.|\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  forward_result = fprop(sparse_indices, values, dense_shape, embedding_name, \n",
    "                         bp_trigger, output_type, is_training=True)\n",
    "  ```\n",
    "  \n",
    "This function can be used to do forward propagation for `distributed` and `localized` embedding layers. It will use all input keys that are stored in the SparseTensor format as its input, and will convert those keys to the CSR format within this function. Therefore, its performance is not very satisfying.\n",
    "  \n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|sparse_indices| A 2-D int64 tensor of shape [N, 3], which specifies the indices of the elements in the sparse tensor that contain valid values. And `N` represents how many valid values in the corresponding dense tensor, 3 represent valid values' [batch_idx, slot_idx, nnz_idx].|\n",
    "|values| A 1-D tensor of type specified in `init().key_type` ans shape [N], which supplies the valid values for each element in `sparse_indices`.|\n",
    "|dense_shape| A 1-D int64 tensor of shape [3], which specifies the dense_shape: `[batch_size, slot_num, max_nnz]` of the sparse tensor.|\n",
    "| embedding_name| tf.Tensor with `dtype=tf.string`, use which embedding layer to do forward propagation.|\n",
    "| bp_trigger| tf.Variable(dtype=tf.float32), used to automatically trigger back propagation of the embedding layer.|\n",
    "| output_type| should be the same with `init().value_type`.|\n",
    "| is_training| bool, specify whether use `training` resources or `evaluation` resources.|\n",
    "\n",
    "|Returns||\n",
    "|:----| :---- |\n",
    "|forward_result| tf.Tensor with `dtype=output_type`. Forward propagation results.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  forward_result = fprop_v3(embedding_name, row_offsets, value_tensors, nnz_array, \n",
    "                            bp_trigger, output_shape, is_training=True)\n",
    "  ```\n",
    "  \n",
    "This function can be used to do forward propagation for `distributed` and `localized` embedding layers. Its inputs has been previously converted to the CSR format. Therefore, no conversion will be conducted within this function. For example, if the `embedding_plugin` uses four GPUs to perform the computation, then four CSR sparse matrices will be needed for its inputs. Addtionally, four row_offsets are stacked together to form a single tensor and value_tensors.\n",
    "\n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, use which embedding layer to do forward propgation.|\n",
    "|row_offsets| 2-D matrix with shape `[gpu_count, batch_size * slot_num + 1]`, `dtype=tf.int64`. Each row in this tensor denotes a CSR `row_offsets` for one GPU.|\n",
    "|value_tensors| 2-D matrix with shape `[gpu_count, keys_nums_in_a_batch]`, `dtype=tf.int64`. Each row in this tensor denotes a CSR `values` for one GPU.|\n",
    "|nnz_array| 1-D tensor with `dtype=tf.int64`, its length is equal to `gpu_count`, and each value denotes how many valid input keys in one CSR sparse matrix.| \n",
    "| bp_trigger| tf.Variable(dtype=tf.float32), used to automatically trigger back propagation of the embedding layer.|\n",
    "| output_shape| 1-D tensor, and its value should be `[batch_size, slot_num, embedding_vec_size]`.|\n",
    "| is_training| bool, specify whether use `training` resources or `evaluation` resources.|\n",
    "\n",
    "|Returns||\n",
    "|:----| :---- |\n",
    "|forward_result| tf.Tensor with `dtype=init().value_type`. Forward propagation result.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  save(embedding_name, save_name)\n",
    "  ```\n",
    "\n",
    "This function is used to save the `embedding_plugin` parameters in the file.\n",
    "\n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, save which embedding layer's parameter to file.|\n",
    "|save_name| string, the name of saved parameters.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ ```python\n",
    "  restore(embedding_name, file_name)\n",
    "  ```\n",
    "  \n",
    "This function is used to restore the `embedding_plugin` parameters from file.\n",
    "\n",
    "|Args||\n",
    "|:----| :---- |\n",
    "|embedding_name| tf.Tensor with `dtype=tf.string`, restore parameters for which embedding layer.|\n",
    "|file_name| string, restore paramters from this file. |"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}