notebooks/introduction_to_tensorflow/solutions/4_keras_functional_api.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introducing the Keras Functional API\n",
    "\n",
    "**Learning Objectives**\n",
    "  1. Understand embeddings and how to create them with the feature column API\n",
    "  1. Understand Deep and Wide models and when to use them\n",
    "  1. Understand the Keras functional API and how to build a deep and wide model with it\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In the last notebook, we learned about the Keras Sequential API. The [Keras Functional API](https://www.tensorflow.org/guide/keras#functional_api) provides an alternate way of building models which is more flexible. With the Functional API, we can build models with more complex topologies, multiple input or output layers, shared layers or non-sequential data flows (e.g. residual layers).\n",
    "\n",
    "In this notebook we'll use what we learned about feature columns to build a Wide & Deep model. Recall, that the idea behind Wide & Deep models is to join the two methods of learning through memorization and generalization by making a wide linear model and a deep learning model to accommodate both. You can have a look at the original research paper here: [Wide & Deep Learning for Recommender Systems](https://arxiv.org/abs/1606.07792).\n",
    "\n",
    "<img src='assets/wide_deep.png' width='80%'>\n",
    "<sup>(image: https://ai.googleblog.com/2016/06/wide-deep-learning-better-together-with.html)</sup>\n",
    "\n",
    "The Wide part of the model is associated with the memory element. In this case, we train a linear model with a wide set of crossed features and learn the correlation of this related data with the assigned label. The Deep part of the model is associated with the generalization element where we use embedding vectors for features. The best embeddings are then learned through the training process. While both of these methods can work well alone, Wide & Deep models excel by combining these techniques together. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start by importing the necessary libraries for this lab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "from matplotlib import pyplot as plt\n",
    "from tensorflow import feature_column as fc\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.callbacks import TensorBoard\n",
    "from tensorflow.keras.layers import Dense, DenseFeatures, Input, concatenate\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load raw data \n",
    "\n",
    "We will use the taxifare dataset, using the CSV files that we created in the first notebook of this sequence. Those files have been saved into `../data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l ../data/*.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use tf.data to read the CSV files\n",
    "\n",
    "We wrote these functions for reading data from the csv files above in the [previous notebook](2_dataset_api.ipynb). For this lab we will also include some additional engineered features in our model. In particular, we will compute the difference in latitude and longitude, as well as the Euclidean distance between the pick-up and drop-off locations. We can accomplish this by adding these new features to the features dictionary with the function `add_engineered_features` below. \n",
    "\n",
    "Note that we include a call to this function when collecting our features dict and labels in the `features_and_labels` function below as well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CSV_COLUMNS = [\n",
    "    \"fare_amount\",\n",
    "    \"pickup_datetime\",\n",
    "    \"pickup_longitude\",\n",
    "    \"pickup_latitude\",\n",
    "    \"dropoff_longitude\",\n",
    "    \"dropoff_latitude\",\n",
    "    \"passenger_count\",\n",
    "    \"key\",\n",
    "]\n",
    "LABEL_COLUMN = \"fare_amount\"\n",
    "DEFAULTS = [[0.0], [\"na\"], [0.0], [0.0], [0.0], [0.0], [0.0], [\"na\"]]\n",
    "UNWANTED_COLS = [\"pickup_datetime\", \"key\"]\n",
    "\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "    features = row_data\n",
    "\n",
    "    for unwanted_col in UNWANTED_COLS:\n",
    "        features.pop(unwanted_col)\n",
    "\n",
    "    return features, label\n",
    "\n",
    "\n",
    "def create_dataset(pattern, batch_size=1, mode=\"eval\"):\n",
    "    dataset = tf.data.experimental.make_csv_dataset(\n",
    "        pattern, batch_size, CSV_COLUMNS, DEFAULTS\n",
    "    )\n",
    "\n",
    "    dataset = dataset.map(features_and_labels)\n",
    "\n",
    "    if mode == \"train\":\n",
    "        dataset = dataset.shuffle(buffer_size=1000).repeat()\n",
    "\n",
    "    # take advantage of multi-threading; 1=AUTOTUNE\n",
    "    dataset = dataset.prefetch(1)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feature columns for Wide and Deep model\n",
    "\n",
    "For the Wide columns, we will create feature columns of crossed features. To do this, we'll create a collection of Tensorflow feature columns to pass to the `tf.feature_column.crossed_column` constructor. The Deep columns will consist of numeric columns and the embedding columns we want to create. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Bucketize latitudes and longitudes\n",
    "NBUCKETS = 16\n",
    "latbuckets = np.linspace(start=38.0, stop=42.0, num=NBUCKETS).tolist()\n",
    "lonbuckets = np.linspace(start=-76.0, stop=-72.0, num=NBUCKETS).tolist()\n",
    "\n",
    "fc_bucketized_plat = fc.bucketized_column(\n",
    "    source_column=fc.numeric_column(\"pickup_latitude\"), boundaries=latbuckets\n",
    ")\n",
    "fc_bucketized_plon = fc.bucketized_column(\n",
    "    source_column=fc.numeric_column(\"pickup_longitude\"), boundaries=lonbuckets\n",
    ")\n",
    "fc_bucketized_dlat = fc.bucketized_column(\n",
    "    source_column=fc.numeric_column(\"dropoff_latitude\"), boundaries=latbuckets\n",
    ")\n",
    "fc_bucketized_dlon = fc.bucketized_column(\n",
    "    source_column=fc.numeric_column(\"dropoff_longitude\"), boundaries=lonbuckets\n",
    ")\n",
    "\n",
    "# 2. Cross features for locations\n",
    "fc_crossed_dloc = fc.crossed_column(\n",
    "    [fc_bucketized_dlat, fc_bucketized_dlon],\n",
    "    hash_bucket_size=NBUCKETS * NBUCKETS,\n",
    ")\n",
    "fc_crossed_ploc = fc.crossed_column(\n",
    "    [fc_bucketized_plat, fc_bucketized_plon],\n",
    "    hash_bucket_size=NBUCKETS * NBUCKETS,\n",
    ")\n",
    "fc_crossed_pd_pair = fc.crossed_column(\n",
    "    [fc_crossed_dloc, fc_crossed_ploc], hash_bucket_size=NBUCKETS**4\n",
    ")\n",
    "\n",
    "# 3. Create embedding columns for the crossed columns\n",
    "fc_pd_pair = fc.embedding_column(\n",
    "    categorical_column=fc_crossed_pd_pair, dimension=3\n",
    ")\n",
    "fc_dloc = fc.embedding_column(categorical_column=fc_crossed_dloc, dimension=3)\n",
    "fc_ploc = fc.embedding_column(categorical_column=fc_crossed_ploc, dimension=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gather list of feature columns\n",
    "\n",
    "Next we gather the list of wide and deep feature columns we'll pass to our Wide & Deep model in Tensorflow. Recall, wide columns are sparse, have linear relationship with the output while continuous columns are deep, have a complex relationship with the output. We will use our previously bucketized columns to collect crossed feature columns and sparse feature columns for our wide columns, and embedding feature columns and numeric features columns for the deep columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wide_columns = [\n",
    "    # One-hot encoded feature crosses\n",
    "    fc.indicator_column(fc_crossed_dloc),\n",
    "    fc.indicator_column(fc_crossed_ploc),\n",
    "    fc.indicator_column(fc_crossed_pd_pair),\n",
    "]\n",
    "\n",
    "deep_columns = [\n",
    "    # Embedding_column to \"group\" together ...\n",
    "    fc.embedding_column(fc_crossed_pd_pair, dimension=10),\n",
    "    # Numeric columns\n",
    "    fc.numeric_column(\"pickup_latitude\"),\n",
    "    fc.numeric_column(\"pickup_longitude\"),\n",
    "    fc.numeric_column(\"dropoff_longitude\"),\n",
    "    fc.numeric_column(\"dropoff_latitude\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build a Wide and Deep model in Keras\n",
    "\n",
    "To build a wide-and-deep network, we connect the sparse (i.e. wide) features directly to the output node, but pass the dense (i.e. deep) features through a set of fully connected layers. Here’s that model architecture looks using the Functional API.\n",
    "\n",
    "First, we'll create our input columns using [tf.keras.layers.Input](https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Input)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT_COLS = [\n",
    "    \"pickup_longitude\",\n",
    "    \"pickup_latitude\",\n",
    "    \"dropoff_longitude\",\n",
    "    \"dropoff_latitude\",\n",
    "    \"passenger_count\",\n",
    "]\n",
    "\n",
    "inputs = {\n",
    "    colname: Input(name=colname, shape=(), dtype=\"float32\")\n",
    "    for colname in INPUT_COLS\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we'll define our custom RMSE evaluation metric and build our wide and deep model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmse(y_true, y_pred):\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))\n",
    "\n",
    "\n",
    "def build_model(dnn_hidden_units):\n",
    "    # Create the deep part of model\n",
    "    deep = DenseFeatures(deep_columns, name=\"deep_inputs\")(inputs)\n",
    "    for num_nodes in dnn_hidden_units:\n",
    "        deep = Dense(num_nodes, activation=\"relu\")(deep)\n",
    "\n",
    "    # Create the wide part of model\n",
    "    wide = DenseFeatures(wide_columns, name=\"wide_inputs\")(inputs)\n",
    "\n",
    "    # Combine deep and wide parts of the model\n",
    "    combined = concatenate(inputs=[deep, wide], name=\"combined\")\n",
    "\n",
    "    # Map the combined outputs into a single prediction value\n",
    "    output = Dense(units=1, activation=None, name=\"prediction\")(combined)\n",
    "\n",
    "    # Finalize the model\n",
    "    model = Model(inputs=list(inputs.values()), outputs=output)\n",
    "\n",
    "    # Compile the keras model\n",
    "    model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[rmse, \"mse\"])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we can call the `build_model` to create the model. Here we'll have two hidden layers, each with 10 neurons, for the deep part of our model. We can also use `plot_model` to see a diagram of the model we've created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HIDDEN_UNITS = [10, 10]\n",
    "\n",
    "model = build_model(dnn_hidden_units=HIDDEN_UNITS)\n",
    "\n",
    "tf.keras.utils.plot_model(model, show_shapes=False, rankdir=\"LR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll set up our training variables, create our datasets for training and validation, and train our model.\n",
    "\n",
    "(We refer you the the blog post [ML Design Pattern #3: Virtual Epochs](https://medium.com/google-cloud/ml-design-pattern-3-virtual-epochs-f842296de730) for further details on why express the training in terms of `NUM_TRAIN_EXAMPLES` and `NUM_EVALS` and why, in this training code, the number of epochs is really equal to the number of evaluations we perform.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1000\n",
    "NUM_TRAIN_EXAMPLES = 10000 * 5  # training dataset will repeat, wrap around\n",
    "NUM_EVALS = 50  # how many times to evaluate\n",
    "NUM_EVAL_EXAMPLES = 10000  # enough to get a reasonable sample\n",
    "\n",
    "trainds = create_dataset(\n",
    "    pattern=\"../data/taxi-train*\", batch_size=BATCH_SIZE, mode=\"train\"\n",
    ")\n",
    "\n",
    "evalds = create_dataset(\n",
    "    pattern=\"../data/taxi-valid*\", batch_size=BATCH_SIZE, mode=\"eval\"\n",
    ").take(NUM_EVAL_EXAMPLES // 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "steps_per_epoch = NUM_TRAIN_EXAMPLES // (BATCH_SIZE * NUM_EVALS)\n",
    "\n",
    "OUTDIR = \"./taxi_trained\"\n",
    "shutil.rmtree(path=OUTDIR, ignore_errors=True)  # start fresh each time\n",
    "\n",
    "history = model.fit(\n",
    "    x=trainds,\n",
    "    steps_per_epoch=steps_per_epoch,\n",
    "    epochs=NUM_EVALS,\n",
    "    validation_data=evalds,\n",
    "    callbacks=[TensorBoard(OUTDIR)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as before, we can examine the history to see how the RMSE changes through training on the train set and validation set. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RMSE_COLS = [\"rmse\", \"val_rmse\"]\n",
    "\n",
    "pd.DataFrame(history.history)[RMSE_COLS].plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright 2021 Google Inc. 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"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-6.m82",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-6:m82"
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}