<commit_before>{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LAB 5a:  Training Keras model on Vertex AI\n",
    "\n",
    "**Learning Objectives**\n",
    "\n",
    "1. Setup up the environment\n",
    "1. Create trainer module's task.py to hold hyperparameter argparsing code\n",
    "1. Create trainer module's model.py to hold Keras model code\n",
    "1. Run trainer module package locally\n",
    "1. Submit training job to Vertex AI\n",
    "1. Submit hyperparameter tuning job to Vertex AI\n",
    "\n",
    "\n",
    "## Introduction\n",
    "After having testing our training pipeline both locally and in the cloud on a susbset of the data, we can submit another (much larger) training job to the cloud. It is also a good idea to run a hyperparameter tuning job to make sure we have optimized the hyperparameters of our model. \n",
    "\n",
    "In this notebook, we'll be training our Keras model at scale using Vertex AI.\n",
    "\n",
    "In this lab, we will set up the environment, create the trainer module's task.py to hold hyperparameter argparsing code, create the trainer module's model.py to hold Keras model code, run the trainer module package locally, submit a training job to Vertex AI, and submit a hyperparameter tuning job to Vertex AI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hJ7ByvoXzpVI"
   },
   "source": [
    "## Set up environment variables and load necessary libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we will install the `cloudml-hypertune` package on our local machine. This is the package which we will use to report hyperparameter tuning metrics to Vertex AI. Installing the package will allow us to test our trainer package locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from google.cloud import aiplatform\n",
    "\n",
    "except ImportError:\n",
    "    !pip3 install -U google-cloud-aiplatform --user\n",
    "\n",
    "    print(\"Please restart the kernel and re-run the notebook.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import hypertune\n",
    "\n",
    "except ImportError:\n",
    "    !pip3 install -U cloudml-hypertune --user\n",
    "\n",
    "    print(\"Please restart the kernel and re-run the notebook.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the above command resulted in an installation, please restart the notebook kernel and re-run the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import necessary libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set environment variables.\n",
    "\n",
    "Set environment variables so that we can use them throughout the entire lab. We will be using our project name for our bucket, so you only need to change your project and region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "export PROJECT=$(gcloud config list project --format \"value(core.project)\")\n",
    "echo \"Your current GCP Project Name is: \"${PROJECT}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Change these to try this notebook out\n",
    "PROJECT = \"asl-ml-immersion\"  # Replace with your PROJECT\n",
    "BUCKET = PROJECT   # defaults to PROJECT\n",
    "REGION = \"us-central1\"  # Replace with your REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"BUCKET\"] = BUCKET\n",
    "os.environ[\"REGION\"] = REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the bucket if does not exist, and confirm below that the bucket is regional and its region equals to the specified region:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " %%bash\n",
    "if ! gsutil ls | grep -q gs://${BUCKET}/; then\n",
    "    gsutil mb -l ${REGION} gs://${BUCKET}\n",
    "fi\n",
    "gsutil ls -Lb gs://$BUCKET | grep \"gs://\\|Location\"\n",
    "echo $REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gcloud config set project ${PROJECT}\n",
    "gcloud config set ai/region ${REGION}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check data exists\n",
    "\n",
    "Verify that you previously created CSV files we'll be using for training and evaluation. If not, go back to lab [1b_prepare_data_babyweight](../solutions/1b_prepare_data_babyweight.ipynb) to create them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gsutil ls gs://${BUCKET}/babyweight/data/*000000000000.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the [Keras wide-and-deep code](../solutions/4c_keras_wide_and_deep_babyweight.ipynb) working on a subset of the data, we can package the TensorFlow code up as a Python module and train it on Vertex AI.\n",
    "\n",
    "## Train on Vertex AI\n",
    "\n",
    "Training on Vertex AI requires:\n",
    "* Making the code a Python source distribution\n",
    "* Using gcloud to submit the training code to [Vertex AI](https://console.cloud.google.com/vertex-ai)\n",
    "\n",
    "Ensure that the Vertex AI API is enabled by going to this [link](https://console.developers.google.com/apis/library/aiplatform.googleapis.com).\n",
    "\n",
    "### Move code into a Python package\n",
    "\n",
    "A Python package is simply a collection of one or more `.py` files along with an `__init__.py` file to identify the containing directory as a package. The `__init__.py` sometimes contains initialization code but for our purposes an empty file suffices.\n",
    "\n",
    "The bash command `touch` creates an empty file in the specified location, the directory `babyweight` should already exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "mkdir -p babyweight/trainer\n",
    "touch babyweight/trainer/__init__.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then use the `%%writefile` magic to write the contents of the cell below to a file called `task.py` in the `babyweight/trainer` folder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create trainer module's task.py to hold hyperparameter argparsing code.\n",
    "\n",
    "The cell below writes the file `babyweight/trainer/task.py` which sets up our training job. Here is where we determine which parameters of our model to pass as flags during training using the `parser` module. Look at how `batch_size` is passed to the model in the code below. Use this as an example to parse arguements for the following variables\n",
    "- `nnsize` which represents the hidden layer sizes to use for DNN feature columns\n",
    "- `nembeds` which represents the embedding size of a cross of n key real-valued parameters\n",
    "- `train_examples` which represents the number of examples (in thousands) to run the training job\n",
    "- `eval_steps` which represents the positive number of steps for which to evaluate model\n",
    "\n",
    "Be sure to include a default value for the parsed arguments above and specfy the `type` if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile babyweight/trainer/task.py\n",
    "import argparse\n",
    "import json\n",
    "import os\n",
    "\n",
    "from trainer import model\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\n",
    "        \"--train_data_path\",\n",
    "        help=\"GCS location of training data\",\n",
    "        required=True\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_data_path\",\n",
    "        help=\"GCS location of evaluation data\",\n",
    "        required=True\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--output_dir\",\n",
    "        help=\"GCS location to write checkpoints and export models\",\n",
    "        default = os.getenv(\"AIP_MODEL_DIR\")\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--batch_size\",\n",
    "        help=\"Number of examples to compute gradient over.\",\n",
    "        type=int,\n",
    "        default=512\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nnsize\",\n",
    "        help=\"Hidden layer sizes for DNN -- provide space-separated layers\",\n",
    "        default=\"128 32 4\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nembeds\",\n",
    "        help=\"Embedding size of a cross of n key real-valued parameters\",\n",
    "        type=int,\n",
    "        default=3\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_epochs\",\n",
    "        help=\"Number of epochs to train the model.\",\n",
    "        type=int,\n",
    "        default=10\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--train_examples\",\n",
    "        help=\"\"\"Number of examples (in thousands) to run the training job over.\n",
    "        If this is more than actual # of examples available, it cycles through\n",
    "        them. So specifying 1000 here when you have only 100k examples makes\n",
    "        this 10 epochs.\"\"\",\n",
    "        type=int,\n",
    "        default=5000\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_steps\",\n",
    "        help=\"\"\"Positive number of steps for which to evaluate model. Default\n",
    "        to None, which means to evaluate until input_fn raises an end-of-input\n",
    "        exception\"\"\",\n",
    "        type=int,\n",
    "        default=None\n",
    "    )\n",
    "\n",
    "    # Parse all arguments\n",
    "    args = parser.parse_args()\n",
    "    arguments = args.__dict__\n",
    "\n",
    "    # Modify some arguments\n",
    "    arguments[\"train_examples\"] *= 1000\n",
    "\n",
    "    # Run the training job\n",
    "    model.train_and_evaluate(arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the same way we can write to the file `model.py` the model that we developed in the previous notebooks. \n",
    "\n",
    "### Create trainer module's model.py to hold Keras model code.\n",
    "\n",
    "To create our `model.py`, we'll use the code we wrote for the Wide & Deep model. Look back at your [9_keras_wide_and_deep_babyweight](../solutions/9_keras_wide_and_deep_babyweight.ipynb) notebook and copy/paste the necessary code from that notebook into its place in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile babyweight/trainer/model.py\n",
    "import datetime\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import hypertune\n",
    "\n",
    "# Determine CSV, label, and key columns\n",
    "CSV_COLUMNS = [\"weight_pounds\",\n",
    "               \"is_male\",\n",
    "               \"mother_age\",\n",
    "               \"plurality\",\n",
    "               \"gestation_weeks\"]\n",
    "LABEL_COLUMN = \"weight_pounds\"\n",
    "\n",
    "# Set default values for each CSV column.\n",
    "# Treat is_male and plurality as strings.\n",
    "DEFAULTS = [[0.0], [\"null\"], [0.0], [\"null\"], [0.0]]\n",
    "\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    \"\"\"Splits features and labels from feature dictionary.\n",
    "\n",
    "    Args:\n",
    "        row_data: Dictionary of CSV column names and tensor values.\n",
    "    Returns:\n",
    "        Dictionary of feature tensors and label tensor.\n",
    "    \"\"\"\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "\n",
    "    return row_data, label  # features, label\n",
    "\n",
    "\n",
    "def load_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):\n",
    "    \"\"\"Loads dataset using the tf.data API from CSV files.\n",
    "\n",
    "    Args:\n",
    "        pattern: str, file pattern to glob into list of files.\n",
    "        batch_size: int, the number of examples per batch.\n",
    "        mode: tf.estimator.ModeKeys to determine if training or evaluating.\n",
    "    Returns:\n",
    "        `Dataset` object.\n",
    "    \"\"\"\n",
    "    print(\"mode = {}\".format(mode))\n",
    "    # Make a CSV dataset\n",
    "    dataset = tf.data.experimental.make_csv_dataset(\n",
    "        file_pattern=pattern,\n",
    "        batch_size=batch_size,\n",
    "        column_names=CSV_COLUMNS,\n",
    "        column_defaults=DEFAULTS)\n",
    "\n",
    "    # Map dataset to features and label\n",
    "    dataset = dataset.map(map_func=features_and_labels)  # features, label\n",
    "\n",
    "    # Shuffle and repeat for training\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        dataset = dataset.shuffle(buffer_size=1000).repeat()\n",
    "\n",
    "    # Take advantage of multi-threading; 1=AUTOTUNE\n",
    "    dataset = dataset.prefetch(buffer_size=1)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "\n",
    "def create_input_layers():\n",
    "    \"\"\"Creates dictionary of input layers for each feature.\n",
    "\n",
    "    Returns:\n",
    "        Dictionary of `tf.Keras.layers.Input` layers for each feature.\n",
    "    \"\"\"\n",
    "    deep_inputs = {\n",
    "        colname: tf.keras.layers.Input(\n",
    "            name=colname, shape=(), dtype=\"float32\")\n",
    "        for colname in [\"mother_age\", \"gestation_weeks\"]\n",
    "    }\n",
    "\n",
    "    wide_inputs = {\n",
    "        colname: tf.keras.layers.Input(\n",
    "            name=colname, shape=(), dtype=\"string\")\n",
    "        for colname in [\"is_male\", \"plurality\"]\n",
    "    }\n",
    "\n",
    "    inputs = {**wide_inputs, **deep_inputs}\n",
    "\n",
    "    return inputs\n",
    "\n",
    "\n",
    "def categorical_fc(name, values):\n",
    "    \"\"\"Helper function to wrap categorical feature by indicator column.\n",
    "\n",
    "    Args:\n",
    "        name: str, name of feature.\n",
    "        values: list, list of strings of categorical values.\n",
    "    Returns:\n",
    "        Categorical and indicator column of categorical feature.\n",
    "    \"\"\"\n",
    "    cat_column = tf.feature_column.categorical_column_with_vocabulary_list(\n",
    "            key=name, vocabulary_list=values)\n",
    "    ind_column = tf.feature_column.indicator_column(\n",
    "        categorical_column=cat_column)\n",
    "\n",
    "    return cat_column, ind_column\n",
    "\n",
    "\n",
    "def create_feature_columns(nembeds):\n",
    "    \"\"\"Creates wide and deep dictionaries of feature columns from inputs.\n",
    "\n",
    "    Args:\n",
    "        nembeds: int, number of dimensions to embed categorical column down to.\n",
    "    Returns:\n",
    "        Wide and deep dictionaries of feature columns.\n",
    "    \"\"\"\n",
    "    deep_fc = {\n",
    "        colname: tf.feature_column.numeric_column(key=colname)\n",
    "        for colname in [\"mother_age\", \"gestation_weeks\"]\n",
    "    }\n",
    "    wide_fc = {}\n",
    "    is_male, wide_fc[\"is_male\"] = categorical_fc(\n",
    "        \"is_male\", [\"True\", \"False\", \"Unknown\"])\n",
    "    plurality, wide_fc[\"plurality\"] = categorical_fc(\n",
    "        \"plurality\", [\"Single(1)\", \"Twins(2)\", \"Triplets(3)\",\n",
    "                      \"Quadruplets(4)\", \"Quintuplets(5)\", \"Multiple(2+)\"])\n",
    "\n",
    "    # Bucketize the float fields. This makes them wide\n",
    "    age_buckets = tf.feature_column.bucketized_column(\n",
    "        source_column=deep_fc[\"mother_age\"],\n",
    "        boundaries=np.arange(15, 45, 1).tolist())\n",
    "    wide_fc[\"age_buckets\"] = tf.feature_column.indicator_column(\n",
    "        categorical_column=age_buckets)\n",
    "\n",
    "    gestation_buckets = tf.feature_column.bucketized_column(\n",
    "        source_column=deep_fc[\"gestation_weeks\"],\n",
    "        boundaries=np.arange(17, 47, 1).tolist())\n",
    "    wide_fc[\"gestation_buckets\"] = tf.feature_column.indicator_column(\n",
    "        categorical_column=gestation_buckets)\n",
    "\n",
    "    # Cross all the wide columns, have to do the crossing before we one-hot\n",
    "    crossed = tf.feature_column.crossed_column(\n",
    "        keys=[age_buckets, gestation_buckets],\n",
    "        hash_bucket_size=1000)\n",
    "    deep_fc[\"crossed_embeds\"] = tf.feature_column.embedding_column(\n",
    "        categorical_column=crossed, dimension=nembeds)\n",
    "\n",
    "    return wide_fc, deep_fc\n",
    "\n",
    "\n",
    "def get_model_outputs(wide_inputs, deep_inputs, dnn_hidden_units):\n",
    "    \"\"\"Creates model architecture and returns outputs.\n",
    "\n",
    "    Args:\n",
    "        wide_inputs: Dense tensor used as inputs to wide side of model.\n",
    "        deep_inputs: Dense tensor used as inputs to deep side of model.\n",
    "        dnn_hidden_units: List of integers where length is number of hidden\n",
    "            layers and ith element is the number of neurons at ith layer.\n",
    "    Returns:\n",
    "        Dense tensor output from the model.\n",
    "    \"\"\"\n",
    "    # Hidden layers for the deep side\n",
    "    layers = [int(x) for x in dnn_hidden_units.split()]\n",
    "    deep = deep_inputs\n",
    "    for layerno, numnodes in enumerate(layers):\n",
    "        deep = tf.keras.layers.Dense(\n",
    "            units=numnodes,\n",
    "            activation=\"relu\",\n",
    "            name=\"dnn_{}\".format(layerno+1))(deep)\n",
    "    deep_out = deep\n",
    "\n",
    "    # Linear model for the wide side\n",
    "    wide_out = tf.keras.layers.Dense(\n",
    "        units=10, activation=\"relu\", name=\"linear\")(wide_inputs)\n",
    "\n",
    "    # Concatenate the two sides\n",
    "    both = tf.keras.layers.concatenate(\n",
    "        inputs=[deep_out, wide_out], name=\"both\")\n",
    "\n",
    "    # Final output is a linear activation because this is regression\n",
    "    output = tf.keras.layers.Dense(\n",
    "        units=1, activation=\"linear\", name=\"weight\")(both)\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "def rmse(y_true, y_pred):\n",
    "    \"\"\"Calculates RMSE evaluation metric.\n",
    "\n",
    "    Args:\n",
    "        y_true: tensor, true labels.\n",
    "        y_pred: tensor, predicted labels.\n",
    "    Returns:\n",
    "        Tensor with value of RMSE between true and predicted labels.\n",
    "    \"\"\"\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))\n",
    "\n",
    "\n",
    "def build_wide_deep_model(dnn_hidden_units=[64, 32], nembeds=3):\n",
    "    \"\"\"Builds wide and deep model using Keras Functional API.\n",
    "\n",
    "    Returns:\n",
    "        `tf.keras.models.Model` object.\n",
    "    \"\"\"\n",
    "    # Create input layers\n",
    "    inputs = create_input_layers()\n",
    "\n",
    "    # Create feature columns for both wide and deep\n",
    "    wide_fc, deep_fc = create_feature_columns(nembeds)\n",
    "\n",
    "    # The constructor for DenseFeatures takes a list of numeric columns\n",
    "    # The Functional API in Keras requires: LayerConstructor()(inputs)\n",
    "    wide_inputs = tf.keras.layers.DenseFeatures(\n",
    "        feature_columns=wide_fc.values(), name=\"wide_inputs\")(inputs)\n",
    "    deep_inputs = tf.keras.layers.DenseFeatures(\n",
    "        feature_columns=deep_fc.values(), name=\"deep_inputs\")(inputs)\n",
    "\n",
    "    # Get output of model given inputs\n",
    "    output = get_model_outputs(wide_inputs, deep_inputs, dnn_hidden_units)\n",
    "\n",
    "    # Build model and compile it all together\n",
    "    model = tf.keras.models.Model(inputs=inputs, outputs=output)\n",
    "    model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[rmse, \"mse\"])\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "# Instantiate the HyperTune reporting object\n",
    "hpt = hypertune.HyperTune()\n",
    "\n",
    "# Reporting callback\n",
    "class HPTCallback(tf.keras.callbacks.Callback):\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        global hpt\n",
    "        hpt.report_hyperparameter_tuning_metric(\n",
    "            hyperparameter_metric_tag='val_rmse',\n",
    "            metric_value=logs['val_rmse'],\n",
    "            global_step=epoch)\n",
    "\n",
    "        \n",
    "def train_and_evaluate(args):\n",
    "    print(args)\n",
    "    model = build_wide_deep_model(args[\"nnsize\"], args[\"nembeds\"])\n",
    "    print(\"Here is our Wide-and-Deep architecture so far:\\n\")\n",
    "    print(model.summary())\n",
    "\n",
    "    trainds = load_dataset(\n",
    "        args[\"train_data_path\"],\n",
    "        args[\"batch_size\"],\n",
    "        tf.estimator.ModeKeys.TRAIN)\n",
    "\n",
    "    evalds = load_dataset(\n",
    "        args[\"eval_data_path\"], 1000, tf.estimator.ModeKeys.EVAL)\n",
    "    if args[\"eval_steps\"]:\n",
    "        evalds = evalds.take(count=args[\"eval_steps\"])\n",
    "\n",
    "    num_batches = args[\"batch_size\"] * args[\"num_epochs\"]\n",
    "    steps_per_epoch = args[\"train_examples\"] // num_batches\n",
    "\n",
    "    checkpoint_path = os.path.join(args[\"output_dir\"], \"checkpoints/babyweight\")\n",<commit_msg>remove debug print<commit_after>{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LAB 5a:  Training Keras model on Vertex AI\n",
    "\n",
    "**Learning Objectives**\n",
    "\n",
    "1. Setup up the environment\n",
    "1. Create trainer module's task.py to hold hyperparameter argparsing code\n",
    "1. Create trainer module's model.py to hold Keras model code\n",
    "1. Run trainer module package locally\n",
    "1. Submit training job to Vertex AI\n",
    "1. Submit hyperparameter tuning job to Vertex AI\n",
    "\n",
    "\n",
    "## Introduction\n",
    "After having testing our training pipeline both locally and in the cloud on a susbset of the data, we can submit another (much larger) training job to the cloud. It is also a good idea to run a hyperparameter tuning job to make sure we have optimized the hyperparameters of our model. \n",
    "\n",
    "In this notebook, we'll be training our Keras model at scale using Vertex AI.\n",
    "\n",
    "In this lab, we will set up the environment, create the trainer module's task.py to hold hyperparameter argparsing code, create the trainer module's model.py to hold Keras model code, run the trainer module package locally, submit a training job to Vertex AI, and submit a hyperparameter tuning job to Vertex AI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hJ7ByvoXzpVI"
   },
   "source": [
    "## Set up environment variables and load necessary libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we will install the `cloudml-hypertune` package on our local machine. This is the package which we will use to report hyperparameter tuning metrics to Vertex AI. Installing the package will allow us to test our trainer package locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from google.cloud import aiplatform\n",
    "\n",
    "except ImportError:\n",
    "    !pip3 install -U google-cloud-aiplatform --user\n",
    "\n",
    "    print(\"Please restart the kernel and re-run the notebook.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import hypertune\n",
    "\n",
    "except ImportError:\n",
    "    !pip3 install -U cloudml-hypertune --user\n",
    "\n",
    "    print(\"Please restart the kernel and re-run the notebook.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the above command resulted in an installation, please restart the notebook kernel and re-run the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import necessary libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set environment variables.\n",
    "\n",
    "Set environment variables so that we can use them throughout the entire lab. We will be using our project name for our bucket, so you only need to change your project and region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "export PROJECT=$(gcloud config list project --format \"value(core.project)\")\n",
    "echo \"Your current GCP Project Name is: \"${PROJECT}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Change these to try this notebook out\n",
    "PROJECT = \"asl-ml-immersion\"  # Replace with your PROJECT\n",
    "BUCKET = PROJECT   # defaults to PROJECT\n",
    "REGION = \"us-central1\"  # Replace with your REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"PROJECT\"] = PROJECT\n",
    "os.environ[\"BUCKET\"] = BUCKET\n",
    "os.environ[\"REGION\"] = REGION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the bucket if does not exist, and confirm below that the bucket is regional and its region equals to the specified region:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " %%bash\n",
    "if ! gsutil ls | grep -q gs://${BUCKET}/; then\n",
    "    gsutil mb -l ${REGION} gs://${BUCKET}\n",
    "fi\n",
    "gsutil ls -Lb gs://$BUCKET | grep \"gs://\\|Location\"\n",
    "echo $REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gcloud config set project ${PROJECT}\n",
    "gcloud config set ai/region ${REGION}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check data exists\n",
    "\n",
    "Verify that you previously created CSV files we'll be using for training and evaluation. If not, go back to lab [1b_prepare_data_babyweight](../solutions/1b_prepare_data_babyweight.ipynb) to create them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "gsutil ls gs://${BUCKET}/babyweight/data/*000000000000.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the [Keras wide-and-deep code](../solutions/4c_keras_wide_and_deep_babyweight.ipynb) working on a subset of the data, we can package the TensorFlow code up as a Python module and train it on Vertex AI.\n",
    "\n",
    "## Train on Vertex AI\n",
    "\n",
    "Training on Vertex AI requires:\n",
    "* Making the code a Python source distribution\n",
    "* Using gcloud to submit the training code to [Vertex AI](https://console.cloud.google.com/vertex-ai)\n",
    "\n",
    "Ensure that the Vertex AI API is enabled by going to this [link](https://console.developers.google.com/apis/library/aiplatform.googleapis.com).\n",
    "\n",
    "### Move code into a Python package\n",
    "\n",
    "A Python package is simply a collection of one or more `.py` files along with an `__init__.py` file to identify the containing directory as a package. The `__init__.py` sometimes contains initialization code but for our purposes an empty file suffices.\n",
    "\n",
    "The bash command `touch` creates an empty file in the specified location, the directory `babyweight` should already exist."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "mkdir -p babyweight/trainer\n",
    "touch babyweight/trainer/__init__.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then use the `%%writefile` magic to write the contents of the cell below to a file called `task.py` in the `babyweight/trainer` folder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create trainer module's task.py to hold hyperparameter argparsing code.\n",
    "\n",
    "The cell below writes the file `babyweight/trainer/task.py` which sets up our training job. Here is where we determine which parameters of our model to pass as flags during training using the `parser` module. Look at how `batch_size` is passed to the model in the code below. Use this as an example to parse arguements for the following variables\n",
    "- `nnsize` which represents the hidden layer sizes to use for DNN feature columns\n",
    "- `nembeds` which represents the embedding size of a cross of n key real-valued parameters\n",
    "- `train_examples` which represents the number of examples (in thousands) to run the training job\n",
    "- `eval_steps` which represents the positive number of steps for which to evaluate model\n",
    "\n",
    "Be sure to include a default value for the parsed arguments above and specfy the `type` if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile babyweight/trainer/task.py\n",
    "import argparse\n",
    "import json\n",
    "import os\n",
    "\n",
    "from trainer import model\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\n",
    "        \"--train_data_path\",\n",
    "        help=\"GCS location of training data\",\n",
    "        required=True\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_data_path\",\n",
    "        help=\"GCS location of evaluation data\",\n",
    "        required=True\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--output_dir\",\n",
    "        help=\"GCS location to write checkpoints and export models\",\n",
    "        default = os.getenv(\"AIP_MODEL_DIR\")\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--batch_size\",\n",
    "        help=\"Number of examples to compute gradient over.\",\n",
    "        type=int,\n",
    "        default=512\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nnsize\",\n",
    "        help=\"Hidden layer sizes for DNN -- provide space-separated layers\",\n",
    "        default=\"128 32 4\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nembeds\",\n",
    "        help=\"Embedding size of a cross of n key real-valued parameters\",\n",
    "        type=int,\n",
    "        default=3\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--num_epochs\",\n",
    "        help=\"Number of epochs to train the model.\",\n",
    "        type=int,\n",
    "        default=10\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--train_examples\",\n",
    "        help=\"\"\"Number of examples (in thousands) to run the training job over.\n",
    "        If this is more than actual # of examples available, it cycles through\n",
    "        them. So specifying 1000 here when you have only 100k examples makes\n",
    "        this 10 epochs.\"\"\",\n",
    "        type=int,\n",
    "        default=5000\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--eval_steps\",\n",
    "        help=\"\"\"Positive number of steps for which to evaluate model. Default\n",
    "        to None, which means to evaluate until input_fn raises an end-of-input\n",
    "        exception\"\"\",\n",
    "        type=int,\n",
    "        default=None\n",
    "    )\n",
    "\n",
    "    # Parse all arguments\n",
    "    args = parser.parse_args()\n",
    "    arguments = args.__dict__\n",
    "\n",
    "    # Modify some arguments\n",
    "    arguments[\"train_examples\"] *= 1000\n",
    "\n",
    "    # Run the training job\n",
    "    model.train_and_evaluate(arguments)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the same way we can write to the file `model.py` the model that we developed in the previous notebooks. \n",
    "\n",
    "### Create trainer module's model.py to hold Keras model code.\n",
    "\n",
    "To create our `model.py`, we'll use the code we wrote for the Wide & Deep model. Look back at your [9_keras_wide_and_deep_babyweight](../solutions/9_keras_wide_and_deep_babyweight.ipynb) notebook and copy/paste the necessary code from that notebook into its place in the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile babyweight/trainer/model.py\n",
    "import datetime\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import hypertune\n",
    "\n",
    "# Determine CSV, label, and key columns\n",
    "CSV_COLUMNS = [\"weight_pounds\",\n",
    "               \"is_male\",\n",
    "               \"mother_age\",\n",
    "               \"plurality\",\n",
    "               \"gestation_weeks\"]\n",
    "LABEL_COLUMN = \"weight_pounds\"\n",
    "\n",
    "# Set default values for each CSV column.\n",
    "# Treat is_male and plurality as strings.\n",
    "DEFAULTS = [[0.0], [\"null\"], [0.0], [\"null\"], [0.0]]\n",
    "\n",
    "\n",
    "def features_and_labels(row_data):\n",
    "    \"\"\"Splits features and labels from feature dictionary.\n",
    "\n",
    "    Args:\n",
    "        row_data: Dictionary of CSV column names and tensor values.\n",
    "    Returns:\n",
    "        Dictionary of feature tensors and label tensor.\n",
    "    \"\"\"\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "\n",
    "    return row_data, label  # features, label\n",
    "\n",
    "\n",
    "def load_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):\n",
    "    \"\"\"Loads dataset using the tf.data API from CSV files.\n",
    "\n",
    "    Args:\n",
    "        pattern: str, file pattern to glob into list of files.\n",
    "        batch_size: int, the number of examples per batch.\n",
    "        mode: tf.estimator.ModeKeys to determine if training or evaluating.\n",
    "    Returns:\n",
    "        `Dataset` object.\n",
    "    \"\"\"\n",
    "    print(\"mode = {}\".format(mode))\n",
    "    # Make a CSV dataset\n",
    "    dataset = tf.data.experimental.make_csv_dataset(\n",
    "        file_pattern=pattern,\n",
    "        batch_size=batch_size,\n",
    "        column_names=CSV_COLUMNS,\n",
    "        column_defaults=DEFAULTS)\n",
    "\n",
    "    # Map dataset to features and label\n",
    "    dataset = dataset.map(map_func=features_and_labels)  # features, label\n",
    "\n",
    "    # Shuffle and repeat for training\n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        dataset = dataset.shuffle(buffer_size=1000).repeat()\n",
    "\n",
    "    # Take advantage of multi-threading; 1=AUTOTUNE\n",
    "    dataset = dataset.prefetch(buffer_size=1)\n",
    "\n",
    "    return dataset\n",
    "\n",
    "\n",
    "def create_input_layers():\n",
    "    \"\"\"Creates dictionary of input layers for each feature.\n",
    "\n",
    "    Returns:\n",
    "        Dictionary of `tf.Keras.layers.Input` layers for each feature.\n",
    "    \"\"\"\n",
    "    deep_inputs = {\n",
    "        colname: tf.keras.layers.Input(\n",
    "            name=colname, shape=(), dtype=\"float32\")\n",
    "        for colname in [\"mother_age\", \"gestation_weeks\"]\n",
    "    }\n",
    "\n",
    "    wide_inputs = {\n",
    "        colname: tf.keras.layers.Input(\n",
    "            name=colname, shape=(), dtype=\"string\")\n",
    "        for colname in [\"is_male\", \"plurality\"]\n",
    "    }\n",
    "\n",
    "    inputs = {**wide_inputs, **deep_inputs}\n",
    "\n",
    "    return inputs\n",
    "\n",
    "\n",
    "def categorical_fc(name, values):\n",
    "    \"\"\"Helper function to wrap categorical feature by indicator column.\n",
    "\n",
    "    Args:\n",
    "        name: str, name of feature.\n",
    "        values: list, list of strings of categorical values.\n",
    "    Returns:\n",
    "        Categorical and indicator column of categorical feature.\n",
    "    \"\"\"\n",
    "    cat_column = tf.feature_column.categorical_column_with_vocabulary_list(\n",
    "            key=name, vocabulary_list=values)\n",
    "    ind_column = tf.feature_column.indicator_column(\n",
    "        categorical_column=cat_column)\n",
    "\n",
    "    return cat_column, ind_column\n",
    "\n",
    "\n",
    "def create_feature_columns(nembeds):\n",
    "    \"\"\"Creates wide and deep dictionaries of feature columns from inputs.\n",
    "\n",
    "    Args:\n",
    "        nembeds: int, number of dimensions to embed categorical column down to.\n",
    "    Returns:\n",
    "        Wide and deep dictionaries of feature columns.\n",
    "    \"\"\"\n",
    "    deep_fc = {\n",
    "        colname: tf.feature_column.numeric_column(key=colname)\n",
    "        for colname in [\"mother_age\", \"gestation_weeks\"]\n",
    "    }\n",
    "    wide_fc = {}\n",
    "    is_male, wide_fc[\"is_male\"] = categorical_fc(\n",
    "        \"is_male\", [\"True\", \"False\", \"Unknown\"])\n",
    "    plurality, wide_fc[\"plurality\"] = categorical_fc(\n",
    "        \"plurality\", [\"Single(1)\", \"Twins(2)\", \"Triplets(3)\",\n",
    "                      \"Quadruplets(4)\", \"Quintuplets(5)\", \"Multiple(2+)\"])\n",
    "\n",
    "    # Bucketize the float fields. This makes them wide\n",
    "    age_buckets = tf.feature_column.bucketized_column(\n",
    "        source_column=deep_fc[\"mother_age\"],\n",
    "        boundaries=np.arange(15, 45, 1).tolist())\n",
    "    wide_fc[\"age_buckets\"] = tf.feature_column.indicator_column(\n",
    "        categorical_column=age_buckets)\n",
    "\n",
    "    gestation_buckets = tf.feature_column.bucketized_column(\n",
    "        source_column=deep_fc[\"gestation_weeks\"],\n",
    "        boundaries=np.arange(17, 47, 1).tolist())\n",
    "    wide_fc[\"gestation_buckets\"] = tf.feature_column.indicator_column(\n",
    "        categorical_column=gestation_buckets)\n",
    "\n",
    "    # Cross all the wide columns, have to do the crossing before we one-hot\n",
    "    crossed = tf.feature_column.crossed_column(\n",
    "        keys=[age_buckets, gestation_buckets],\n",
    "        hash_bucket_size=1000)\n",
    "    deep_fc[\"crossed_embeds\"] = tf.feature_column.embedding_column(\n",
    "        categorical_column=crossed, dimension=nembeds)\n",
    "\n",
    "    return wide_fc, deep_fc\n",
    "\n",
    "\n",
    "def get_model_outputs(wide_inputs, deep_inputs, dnn_hidden_units):\n",
    "    \"\"\"Creates model architecture and returns outputs.\n",
    "\n",
    "    Args:\n",
    "        wide_inputs: Dense tensor used as inputs to wide side of model.\n",
    "        deep_inputs: Dense tensor used as inputs to deep side of model.\n",
    "        dnn_hidden_units: List of integers where length is number of hidden\n",
    "            layers and ith element is the number of neurons at ith layer.\n",
    "    Returns:\n",
    "        Dense tensor output from the model.\n",
    "    \"\"\"\n",
    "    # Hidden layers for the deep side\n",
    "    layers = [int(x) for x in dnn_hidden_units.split()]\n",
    "    deep = deep_inputs\n",
    "    for layerno, numnodes in enumerate(layers):\n",
    "        deep = tf.keras.layers.Dense(\n",
    "            units=numnodes,\n",
    "            activation=\"relu\",\n",
    "            name=\"dnn_{}\".format(layerno+1))(deep)\n",
    "    deep_out = deep\n",
    "\n",
    "    # Linear model for the wide side\n",
    "    wide_out = tf.keras.layers.Dense(\n",
    "        units=10, activation=\"relu\", name=\"linear\")(wide_inputs)\n",
    "\n",
    "    # Concatenate the two sides\n",
    "    both = tf.keras.layers.concatenate(\n",
    "        inputs=[deep_out, wide_out], name=\"both\")\n",
    "\n",
    "    # Final output is a linear activation because this is regression\n",
    "    output = tf.keras.layers.Dense(\n",
    "        units=1, activation=\"linear\", name=\"weight\")(both)\n",
    "\n",
    "    return output\n",
    "\n",
    "\n",
    "def rmse(y_true, y_pred):\n",
    "    \"\"\"Calculates RMSE evaluation metric.\n",
    "\n",
    "    Args:\n",
    "        y_true: tensor, true labels.\n",
    "        y_pred: tensor, predicted labels.\n",
    "    Returns:\n",
    "        Tensor with value of RMSE between true and predicted labels.\n",
    "    \"\"\"\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true)))\n",
    "\n",
    "\n",
    "def build_wide_deep_model(dnn_hidden_units=[64, 32], nembeds=3):\n",
    "    \"\"\"Builds wide and deep model using Keras Functional API.\n",
    "\n",
    "    Returns:\n",
    "        `tf.keras.models.Model` object.\n",
    "    \"\"\"\n",
    "    # Create input layers\n",
    "    inputs = create_input_layers()\n",
    "\n",
    "    # Create feature columns for both wide and deep\n",
    "    wide_fc, deep_fc = create_feature_columns(nembeds)\n",
    "\n",
    "    # The constructor for DenseFeatures takes a list of numeric columns\n",
    "    # The Functional API in Keras requires: LayerConstructor()(inputs)\n",
    "    wide_inputs = tf.keras.layers.DenseFeatures(\n",
    "        feature_columns=wide_fc.values(), name=\"wide_inputs\")(inputs)\n",
    "    deep_inputs = tf.keras.layers.DenseFeatures(\n",
    "        feature_columns=deep_fc.values(), name=\"deep_inputs\")(inputs)\n",
    "\n",
    "    # Get output of model given inputs\n",
    "    output = get_model_outputs(wide_inputs, deep_inputs, dnn_hidden_units)\n",
    "\n",
    "    # Build model and compile it all together\n",
    "    model = tf.keras.models.Model(inputs=inputs, outputs=output)\n",
    "    model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[rmse, \"mse\"])\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "# Instantiate the HyperTune reporting object\n",
    "hpt = hypertune.HyperTune()\n",
    "\n",
    "# Reporting callback\n",
    "class HPTCallback(tf.keras.callbacks.Callback):\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        global hpt\n",
    "        hpt.report_hyperparameter_tuning_metric(\n",
    "            hyperparameter_metric_tag='val_rmse',\n",
    "            metric_value=logs['val_rmse'],\n",
    "            global_step=epoch)\n",
    "\n",
    "        \n",
    "def train_and_evaluate(args):\n",
    "    model = build_wide_deep_model(args[\"nnsize\"], args[\"nembeds\"])\n",
    "    print(\"Here is our Wide-and-Deep architecture so far:\\n\")\n",
    "    print(model.summary())\n",
    "\n",
    "    trainds = load_dataset(\n",
    "        args[\"train_data_path\"],\n",
    "        args[\"batch_size\"],\n",
    "        tf.estimator.ModeKeys.TRAIN)\n",
    "\n",
    "    evalds = load_dataset(\n",
    "        args[\"eval_data_path\"], 1000, tf.estimator.ModeKeys.EVAL)\n",
    "    if args[\"eval_steps\"]:\n",
    "        evalds = evalds.take(count=args[\"eval_steps\"])\n",
    "\n",
    "    num_batches = args[\"batch_size\"] * args[\"num_epochs\"]\n",
    "    steps_per_epoch = args[\"train_examples\"] // num_batches\n",
    "\n",
    "    checkpoint_path = os.path.join(args[\"output_dir\"], \"checkpoints/babyweight\")\n",
{'fraction_non_alphanumeric': 0.17349280607291023, 'fraction_numerical': 0.004634124812428282, 'mean_word_length': 2.316282473472375, 'pattern_counts': {'":': 200, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7155438', 'n_tokens_mistral': 15567, 'n_tokens_neox': 15186, 'n_words': 3830}