documentation/tutorials/intermediate_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title 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",
        "# https://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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8yo62ffS5TF5"
      },
      "source": [
        "# Using text and neural network features\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/decision_forests/tutorials/intermediate_colab\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/decision-forests/blob/main/documentation/tutorials/intermediate_colab.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/decision-forests/blob/main/documentation/tutorials/intermediate_colab.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/decision-forests/documentation/tutorials/intermediate_colab.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://tfhub.dev/google/universal-sentence-encoder/4\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" /\u003eSee TF Hub model\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zrCwCCxhiAL7"
      },
      "source": [
        "Welcome to the **Intermediate Colab** for **TensorFlow Decision Forests** (**TF-DF**).\n",
        "In this colab, you will learn about some more advanced capabilities of **TF-DF**, including how to deal with natural language features.\n",
        "\n",
        "This colab assumes you are familiar with the concepts presented the [Beginner colab](beginner_colab.ipynb), notably about the installation about TF-DF.\n",
        "\n",
        "In this colab, you will:\n",
        "\n",
        "1. Train a Random Forest that consumes text features natively as categorical sets.\n",
        "\n",
        "1. Train a Random Forest that consumes text features using a [TensorFlow Hub](https://www.tensorflow.org/hub) module. In this setting (transfer learning), the module is already pre-trained on a large text corpus.\n",
        "\n",
        "1. Train a Gradient Boosted Decision Trees (GBDT) and a Neural Network together. The GBDT will consume the output of the Neural Network."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rzskapxq7gdo"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mZiInVYfffAb"
      },
      "outputs": [],
      "source": [
        "# Install TensorFlow Dececision Forests\n",
        "!pip install tensorflow_decision_forests\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2EFndCFdoJM5"
      },
      "source": [
        "[Wurlitzer](https://pypi.org/project/wurlitzer/) is needed to display the detailed training logs in Colabs (when using `verbose=2` in the model constructor)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L06XWRdSoLj5"
      },
      "outputs": [],
      "source": [
        "!pip install wurlitzer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7PlfbnxYcPf"
      },
      "source": [
        "Import the necessary libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RsCV2oAS7gC_"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "# Keep using Keras 2\n",
        "os.environ['TF_USE_LEGACY_KERAS'] = '1'\n",
        "\n",
        "import tensorflow_decision_forests as tfdf\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import tensorflow as tf\n",
        "import tf_keras\n",
        "import math"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w2fsI0y5x5i5"
      },
      "source": [
        "The hidden code cell limits the output height in colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jZXB4o6Tlu0i"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "\n",
        "from IPython.core.magic import register_line_magic\n",
        "from IPython.display import Javascript\n",
        "from IPython.display import display as ipy_display\n",
        "\n",
        "# Some of the model training logs can cover the full\n",
        "# screen if not compressed to a smaller viewport.\n",
        "# This magic allows setting a max height for a cell.\n",
        "@register_line_magic\n",
        "def set_cell_height(size):\n",
        "  ipy_display(\n",
        "      Javascript(\"google.colab.output.setIframeHeight(0, true, {maxHeight: \" +\n",
        "                 str(size) + \"})\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M_D4Ft4o65XT"
      },
      "source": [
        "## Use raw text as features\n",
        "\n",
        "TF-DF can consume [categorical-set](https://arxiv.org/pdf/2009.09991.pdf) features natively. Categorical-sets represent text features as bags of words (or n-grams).\n",
        "\n",
        "For example: `\"The little blue dog\" ` → `{\"the\", \"little\", \"blue\", \"dog\"}`\n",
        "\n",
        "In this example, you'll will train a Random Forest on the [Stanford Sentiment Treebank](https://nlp.stanford.edu/sentiment/index.html) (SST) dataset. The objective of this dataset is to classify sentences as carrying a *positive* or *negative* sentiment. You'll will use the binary classification version of the dataset curated in [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/glue#gluesst2).\n",
        "\n",
        "**Note:** Categorical-set features can be expensive to train. In this colab, we will train a small Random Forest with 20 trees."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SgEiFy23j14S"
      },
      "outputs": [],
      "source": [
        "# Install the TensorFlow Datasets package\n",
        "!pip install tensorflow-datasets -U --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uVN-j0E4Q1T3"
      },
      "outputs": [],
      "source": [
        "# Load the dataset\n",
        "import tensorflow_datasets as tfds\n",
        "all_ds = tfds.load(\"glue/sst2\")\n",
        "\n",
        "# Display the first 3 examples of the test fold.\n",
        "for example in all_ds[\"test\"].take(3):\n",
        "  print({attr_name: attr_tensor.numpy() for attr_name, attr_tensor in example.items()})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UHiQUWE2XDYN"
      },
      "source": [
        "The dataset is modified as follows:\n",
        "\n",
        "1. The raw labels are integers in `{-1, 1}`, but the learning algorithm expects positive integer labels e.g. `{0, 1}`. Therefore, the labels are transformed as follows: `new_labels = (original_labels + 1) / 2`.\n",
        "1. A batch-size of 64 is applied to make reading the dataset more efficient.\n",
        "1. The `sentence` attribute needs to be tokenized, i.e. `\"hello world\" -\u003e [\"hello\", \"world\"]`.\n",
        "\n",
        "\n",
        "**Note:** This example doesn't use the `test` split of the dataset as it does not have labels. If `test` split had labels, you could concatenate the `validation` fold into the `train` one (e.g. `all_ds[\"train\"].concatenate(all_ds[\"validation\"])`).\n",
        "\n",
        "**Details:** Some decision forest learning algorithms do not need a validation dataset (e.g. Random Forests) while others do (e.g. Gradient Boosted Trees in some cases). Since each learning algorithm under TF-DF can use validation data differently, TF-DF handles train/validation splits internally. As a result, when you have a training and validation sets, they can always be concatenated as input to the learning algorithm."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yqYDKTKdSPYw"
      },
      "outputs": [],
      "source": [
        "def prepare_dataset(example):\n",
        "  label = (example[\"label\"] + 1) // 2\n",
        "  return {\"sentence\" : tf.strings.split(example[\"sentence\"])}, label\n",
        "\n",
        "train_ds = all_ds[\"train\"].batch(100).map(prepare_dataset)\n",
        "test_ds = all_ds[\"validation\"].batch(100).map(prepare_dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YYkIjROI9w43"
      },
      "source": [
        "Finally, train and evaluate the model as usual. TF-DF  automatically detects multi-valued categorical features as categorical-set.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mpxTtYo39wYZ"
      },
      "outputs": [],
      "source": [
        "%set_cell_height 300\n",
        "\n",
        "# Specify the model.\n",
        "model_1 = tfdf.keras.RandomForestModel(num_trees=30, verbose=2)\n",
        "\n",
        "# Train the model.\n",
        "model_1.fit(x=train_ds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9FMFGzwiHCt"
      },
      "source": [
        "In the previous logs, note that `sentence` is a `CATEGORICAL_SET` feature.\n",
        "\n",
        "The model is evaluated as usual:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cpf-wHl094S1"
      },
      "outputs": [],
      "source": [
        "model_1.compile(metrics=[\"accuracy\"])\n",
        "evaluation = model_1.evaluate(test_ds)\n",
        "\n",
        "print(f\"BinaryCrossentropyloss: {evaluation[0]}\")\n",
        "print(f\"Accuracy: {evaluation[1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YliBX4GtjncQ"
      },
      "source": [
        "The training logs looks are follow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OnTTtBNmjpo7"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "\n",
        "logs = model_1.make_inspector().training_logs()\n",
        "plt.plot([log.num_trees for log in logs], [log.evaluation.accuracy for log in logs])\n",
        "plt.xlabel(\"Number of trees\")\n",
        "plt.ylabel(\"Out-of-bag accuracy\")\n",
        "pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d4qJ0ig3kgic"
      },
      "source": [
        "More trees would probably be beneficial (I am sure of it because I tried :p)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iil_oyOhCNx6"
      },
      "source": [
        "## Use a pretrained text embedding\n",
        "\n",
        "The previous example trained a Random Forest using raw text features. This example will use a pre-trained TF-Hub embedding to convert text features into a dense embedding, and then train a Random Forest on top of it. In this situation, the Random Forest will only \"see\" the numerical output of the embedding (i.e. it will not see the raw text). \n",
        "\n",
        "In this experiment,  will use the [Universal-Sentence-Encoder](https://tfhub.dev/google/universal-sentence-encoder/4). Different pre-trained embeddings might be suited for different types of text (e.g. different language, different task) but also for other type of structured features (e.g. images).\n",
        "\n",
        "**Note:** This embedding is large (1GB) and therefore the final model will be slow to run (compared to classical decision tree inference).\n",
        "\n",
        "The embedding module can be applied in one of two places:\n",
        "\n",
        "1. During the dataset preparation.\n",
        "2. In the pre-processing stage of the model.\n",
        "\n",
        "The second option is often preferable: Packaging the embedding in the model makes the model easier to use (and harder to misuse).\n",
        "\n",
        "First install TF-Hub:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QfYGXim_DskC"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade tensorflow-hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kNSEhJgjEXww"
      },
      "source": [
        "Unlike before, you don't need to tokenize the text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pS5SYqoScbOc"
      },
      "outputs": [],
      "source": [
        "def prepare_dataset(example):\n",
        "  label = (example[\"label\"] + 1) // 2\n",
        "  return {\"sentence\" : example[\"sentence\"]}, label\n",
        "\n",
        "train_ds = all_ds[\"train\"].batch(100).map(prepare_dataset)\n",
        "test_ds = all_ds[\"validation\"].batch(100).map(prepare_dataset)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zHEsd8q_ESpC"
      },
      "outputs": [],
      "source": [
        "%set_cell_height 300\n",
        "\n",
        "import tensorflow_hub as hub\n",
        "# NNLM (https://tfhub.dev/google/nnlm-en-dim128/2) is also a good choice.\n",
        "hub_url = \"https://tfhub.dev/google/universal-sentence-encoder/4\"\n",
        "embedding = hub.KerasLayer(hub_url)\n",
        "\n",
        "sentence = tf_keras.layers.Input(shape=(), name=\"sentence\", dtype=tf.string)\n",
        "embedded_sentence = embedding(sentence)\n",
        "\n",
        "raw_inputs = {\"sentence\": sentence}\n",
        "processed_inputs = {\"embedded_sentence\": embedded_sentence}\n",
        "preprocessor = tf_keras.Model(inputs=raw_inputs, outputs=processed_inputs)\n",
        "\n",
        "model_2 = tfdf.keras.RandomForestModel(\n",
        "    preprocessing=preprocessor,\n",
        "    num_trees=100)\n",
        "\n",
        "model_2.fit(x=train_ds)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xPLoDqiFKY18"
      },
      "outputs": [],
      "source": [
        "model_2.compile(metrics=[\"accuracy\"])\n",
        "evaluation = model_2.evaluate(test_ds)\n",
        "\n",
        "print(f\"BinaryCrossentropyloss: {evaluation[0]}\")\n",
        "print(f\"Accuracy: {evaluation[1]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WPsD3LyaMLHm"
      },
      "source": [
        "Note that categorical sets represent text differently from a dense embedding, so it may be useful to use both strategies jointly."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "37AGJamzboZQ"
      },
      "source": [
        "## Train a decision tree and neural network together\n",
        "\n",
        "The previous example used a pre-trained Neural Network (NN) to \n",
        "process the text features before passing them to the Random Forest. This example will train both the Neural Network and the Random Forest from scratch.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YJIxGwwzMkFl"
      },
      "source": [
        "TF-DF's Decision Forests do not back-propagate gradients ([although this is the subject of ongoing research](https://arxiv.org/abs/2007.14761)). Therefore, the training happens in two stages:\n",
        "\n",
        "1. Train the neural-network as a standard classification task:\n",
        "\n",
        "```\n",
        "example → [Normalize] → [Neural Network*] → [classification head] → prediction\n",
        "*: Training.\n",
        "```\n",
        "\n",
        "2. Replace the Neural Network's head (the last layer and the soft-max) with a Random Forest. Train the Random Forest as usual:\n",
        "\n",
        "```\n",
        "example → [Normalize] → [Neural Network] → [Random Forest*] → prediction\n",
        "*: Training.\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YSIvuAhzbjWO"
      },
      "source": [
        "### Prepare the dataset\n",
        "\n",
        "This example uses the [Palmer's Penguins](https://allisonhorst.github.io/palmerpenguins/articles/intro.html) dataset. See the [Beginner colab](beginner_colab.ipynb) for details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "InUot_K2b3Mz"
      },
      "source": [
        "First, download the raw data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rNyaeCx0b1be"
      },
      "outputs": [],
      "source": [
        "!wget -q https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins.csv -O /tmp/penguins.csv"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pNPZzQekb9z_"
      },
      "source": [
        "Load a dataset into a Pandas Dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9lA3peQ4sa9a"
      },
      "outputs": [],
      "source": [
        "dataset_df = pd.read_csv(\"/tmp/penguins.csv\")\n",
        "\n",
        "# Display the first 3 examples.\n",
        "dataset_df.head(3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v-_SZpRWcAoX"
      },
      "source": [
        "\n",
        "Prepare the dataset for training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rtyi8UoqtzhM"
      },
      "outputs": [],
      "source": [
        "label = \"species\"\n",
        "\n",
        "# Replaces numerical NaN (representing missing values in Pandas Dataframe) with 0s.\n",
        "# ...Neural Nets don't work well with numerical NaNs.\n",
        "for col in dataset_df.columns:\n",
        "  if dataset_df[col].dtype not in [str, object]:\n",
        "    dataset_df[col] = dataset_df[col].fillna(0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GKrW5Yfjso0k"
      },
      "outputs": [],
      "source": [
        "# Split the dataset into a training and testing dataset.\n",
        "\n",
        "def split_dataset(dataset, test_ratio=0.30):\n",
        "  \"\"\"Splits a panda dataframe in two.\"\"\"\n",
        "  test_indices = np.random.rand(len(dataset)) \u003c test_ratio\n",
        "  return dataset[~test_indices], dataset[test_indices]\n",
        "\n",
        "train_ds_pd, test_ds_pd = split_dataset(dataset_df)\n",
        "print(\"{} examples in training, {} examples for testing.\".format(\n",
        "    len(train_ds_pd), len(test_ds_pd)))\n",
        "\n",
        "# Convert the datasets into tensorflow datasets\n",
        "train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(train_ds_pd, label=label)\n",
        "test_ds = tfdf.keras.pd_dataframe_to_tf_dataset(test_ds_pd, label=label)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ore7f6tgcOMh"
      },
      "source": [
        "### Build the models\n",
        "\n",
        "Next create the neural network model using [Keras' functional style](https://www.tensorflow.org/guide/keras/functional). \n",
        "\n",
        "To keep the example simple this model only uses two inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S1Jfe4YteBqY"
      },
      "outputs": [],
      "source": [
        "input_1 = tf_keras.Input(shape=(1,), name=\"bill_length_mm\", dtype=\"float\")\n",
        "input_2 = tf_keras.Input(shape=(1,), name=\"island\", dtype=\"string\")\n",
        "\n",
        "nn_raw_inputs = [input_1, input_2]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZjlvAUNGeDM8"
      },
      "source": [
        "Use [preprocessing layers](https://www.tensorflow.org/guide/keras/preprocessing_layers) to convert the raw inputs to inputs appropriate for the neural network. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Q09Nkp6ei21"
      },
      "outputs": [],
      "source": [
        "# Normalization.\n",
        "Normalization = tf_keras.layers.Normalization\n",
        "CategoryEncoding = tf_keras.layers.CategoryEncoding\n",
        "StringLookup = tf_keras.layers.StringLookup\n",
        "\n",
        "values = train_ds_pd[\"bill_length_mm\"].values[:, tf.newaxis]\n",
        "input_1_normalizer = Normalization()\n",
        "input_1_normalizer.adapt(values)\n",
        "\n",
        "values = train_ds_pd[\"island\"].values\n",
        "input_2_indexer = StringLookup(max_tokens=32)\n",
        "input_2_indexer.adapt(values)\n",
        "\n",
        "input_2_onehot = CategoryEncoding(output_mode=\"binary\", max_tokens=32)\n",
        "\n",
        "normalized_input_1 = input_1_normalizer(input_1)\n",
        "normalized_input_2 = input_2_onehot(input_2_indexer(input_2))\n",
        "\n",
        "nn_processed_inputs = [normalized_input_1, normalized_input_2]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZCoQljyhelau"
      },
      "source": [
        "Build the body of the neural network:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KzocgbYNsH6y"
      },
      "outputs": [],
      "source": [
        "y = tf_keras.layers.Concatenate()(nn_processed_inputs)\n",
        "y = tf_keras.layers.Dense(16, activation=tf.nn.relu6)(y)\n",
        "last_layer = tf_keras.layers.Dense(8, activation=tf.nn.relu, name=\"last\")(y)\n",
        "\n",
        "# \"3\" for the three label classes. If it were a binary classification, the\n",
        "# output dim would be 1.\n",
        "classification_output = tf_keras.layers.Dense(3)(y)\n",
        "\n",
        "nn_model = tf_keras.models.Model(nn_raw_inputs, classification_output)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zPbRKf1CfIrj"
      },
      "source": [
        "This `nn_model` directly produces classification logits. \n",
        "\n",
        "Next create a decision forest model. This will operate on the high level features that the neural network extracts in the last layer before that classification head."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7fnpGNyTuXvH"
      },
      "outputs": [],
      "source": [
        "# To reduce the risk of mistakes, group both the decision forest and the\n",
        "# neural network in a single keras model.\n",
        "nn_without_head = tf_keras.models.Model(inputs=nn_model.inputs, outputs=last_layer)\n",
        "df_and_nn_model = tfdf.keras.RandomForestModel(preprocessing=nn_without_head)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "trq07lvMudlz"
      },
      "source": [
        "### Train and evaluate the models\n",
        "\n",
        "The model will be trained in two stages. First train the neural network with its own classification head:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h4OyUWKiupuF"
      },
      "outputs": [],
      "source": [
        "%set_cell_height 300\n",
        "\n",
        "nn_model.compile(\n",
        "  optimizer=tf_keras.optimizers.Adam(),\n",
        "  loss=tf_keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "  metrics=[\"accuracy\"])\n",
        "\n",
        "nn_model.fit(x=train_ds, validation_data=test_ds, epochs=10)\n",
        "nn_model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N2mgMZOpgMQp"
      },
      "source": [
        "The neural network layers are shared between the two models. So now that the neural network is trained the decision forest model will be fit to the trained output of the neural network layers:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JAc9niXqud7V"
      },
      "outputs": [],
      "source": [
        "%set_cell_height 300\n",
        "\n",
        "df_and_nn_model.fit(x=train_ds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HF8Ru2HSv1a5"
      },
      "source": [
        "Now evaluate the composed model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EPMlcObzuw89"
      },
      "outputs": [],
      "source": [
        "df_and_nn_model.compile(metrics=[\"accuracy\"])\n",
        "print(\"Evaluation:\", df_and_nn_model.evaluate(test_ds))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "awiHEznlv5sI"
      },
      "source": [
        "Compare it to the Neural Network alone:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "--ompWYTvxM-"
      },
      "outputs": [],
      "source": [
        "print(\"Evaluation :\", nn_model.evaluate(test_ds))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "intermediate_colab.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}