Merge branch 'main' into autolog_tf

notebooks/community/ml_ops/stage2/get_started_with_vertex_experiments_autologging_tf.ipynb

@@ -802,9 +802,9 @@
         "id": "ce76826902c0"
       },
       "source": [
-        "### Train the model with Vertex AI Experiments\n",
+        "### Train a TensorFlow Sequential model with Vertex AI Experiments\n",
         "\n",
-        "In the following code, you build, train and evaluate a TensorFlow tabular model. The Python script includes the following calls to integrate `Vertex AI Experiments`:\n",
+        "In the following code, you build, train and evaluate a TensorFlow Sequential tabular model. The Python script includes the following calls to integrate `Vertex AI Experiments`:\n",
         "\n",
         "- command-line arguments: The arguments `experiment` and `run` are used to pass in the experiment and run names for the experiment.\n",
         "- `autologging()`: Initializes the experiment and does the heap injection.\n",
@@ -908,6 +908,116 @@
         "experiment_df.T"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ce76826902c0"
+      },
+      "source": [
+        "### Train a TensorFlow Functional model with Vertex AI Experiments\n",
+        "\n",
+        "In the following code, you build, train and evaluate a TensorFlow Functional tabular model. The Python script includes the following calls to integrate `Vertex AI Experiments`:\n",
+        "\n",
+        "- command-line arguments: The arguments `experiment` and `run` are used to pass in the experiment and run names for the experiment.\n",
+        "- `autologging()`: Initializes the experiment and does the heap injection.\n",
+        "- `aiplatform.start_execution()`: Initializes a context for linking artifacts.\n",
+        "- `aiplatform.end_run()`: Ends the experiment.\n",
+        "\n",
+        "*Note:* The initializer `Model` will be redirected to `VertexTFModel` by heap injection. When subsequent calls are made to the compile(), fit() and evaluate() methods, they will be executed as the corresponding `VertexTFModel` methods."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "427846783ed6"
+      },
+      "outputs": [],
+      "source": [
+        "RUN_NAME = \"run-2\"\n",
+        "\n",
+        "\n",
+        "def make_dataset():\n",
+        "\n",
+        "    # Scaling Boston Housing data features\n",
+        "    def scale(feature):\n",
+        "        max = np.max(feature)\n",
+        "        feature = (feature / max).astype(np.float)\n",
+        "        return feature, max\n",
+        "\n",
+        "    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
+        "        path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
+        "    )\n",
+        "    params = []\n",
+        "\n",
+        "    for _ in range(13):\n",
+        "        x_train[_], max = scale(x_train[_])\n",
+        "        x_test[_], _ = scale(x_test[_])\n",
+        "        params.append(max)\n",
+        "\n",
+        "    return (x_train, y_train), (x_test, y_test)\n",
+        "\n",
+        "\n",
+        "# Build the Keras model\n",
+        "def build_and_compile_dnn_model(lr):\n",
+        "    inputs = tf.keras.Input(shape=(13,))\n",
+        "    x = tf.keras.layers.Dense(128, activation=\"relu\")(inputs)\n",
+        "    x = tf.keras.layers.Dense(128, activation=\"relu\")(x)\n",
+        "    outputs = tf.keras.layers.Dense(1, activation=\"linear\")(x)\n",
+        "\n",
+        "    model = tf.keras.Model(inputs, outputs)\n",
+        "\n",
+        "    model.compile(\n",
+        "        loss=\"mse\",\n",
+        "        optimizer=tf.keras.optimizers.RMSprop(learning_rate=lr),\n",
+        "        metrics=[tf.keras.metrics.RootMeanSquaredError()],\n",
+        "    )\n",
+        "    return model\n",
+        "\n",
+        "\n",
+        "# autologging\n",
+        "autolog(experiment=EXPERIMENT_NAME, run=RUN_NAME)\n",
+        "\n",
+        "with aiplatform.start_execution(\n",
+        "    schema_title=\"system.ContainerExecution\", display_name=\"example_training\"\n",
+        ") as execution:\n",
+        "    BATCH_SIZE = 16\n",
+        "\n",
+        "    model = build_and_compile_dnn_model(lr=0.01)\n",
+        "\n",
+        "    # Train the model\n",
+        "    (x_train, y_train), (x_test, y_test) = make_dataset()\n",
+        "    model.fit(x_train, y_train, epochs=10, batch_size=BATCH_SIZE)\n",
+        "\n",
+        "    model.evaluate(x_test, y_test)\n",
+        "\n",
+        "aiplatform.end_run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5f40912e6500"
+      },
+      "source": [
+        "#### Get the experiment results\n",
+        "\n",
+        "Next, you use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7e9671712230"
+      },
+      "outputs": [],
+      "source": [
+        "experiment_df = aiplatform.get_experiment_df()\n",
+        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
+        "experiment_df.T"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {