feat: log model size

notebooks/community/ml_ops/stage2/get_started_with_vertex_experiments_autologging_tf.ipynb

@@ -82,9 +82,11 @@
         "The steps performed include:\n",
         "\n",
         "- Construct the DIY autologging code.\n",
-        "- Construct training package with call to autologging.\n",
+        "- Construct training package for TensorFlow Sequential model with call to autologging.\n",
         "- Train a model.\n",
         "- View the experiment\n",
+        "- Construct training package for TensorFlow Functional model with call to autologging.\n",
+        "- Compare the experiment runs.\n",
         "- Delete the experiment."
       ]
     },
@@ -94,9 +96,9 @@
         "id": "2739272aae1b"
       },
       "source": [
-        "### Model\n",
+        "### Dataset\n",
         "\n",
-        "The model used for this tutorial is a pretrain TensorFlow model that was trained on the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+        "The dataset used in this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
       ]
     },
     {
@@ -481,7 +483,8 @@
         "    - `compile()`: overridden method of super class. Automatically logs specified hyperparameters and calls the underlying `compile()` method.\n",
         "    - `fit()`: overridden method of super class. Automatically logs specified hyperparameters, calls the underlying `fit()` method, and logs the resulting metrics.\n",
         "    - `evaluate()`: overridden method of super class. Calls the underlying `evaluate()` method, and logs the resulting metrics.\n",
-        "- `VertexTFModel`: A subclass of the tf.keras.Model class."
+        "- `VertexTFModel`: A subclass of the tf.keras.Model class.\n",
+        "- `VertexTFHelper`: A class for common logging methods for both Sequential and Functional models."
       ]
     },
     {
@@ -630,6 +633,8 @@
         "            use_multiprocessing=use_multiprocessing,\n",
         "        )\n",
         "\n",
+        "        TFHelper().model_size(self)\n",
+        "\n",
         "        for key, val in history.history.items():\n",
         "            aiplatform.log_metrics({f\"train.{key}\": val[-1]})\n",
         "        return history\n",
@@ -757,6 +762,8 @@
         "            use_multiprocessing=use_multiprocessing,\n",
         "        )\n",
         "\n",
+        "        TFHelper().model_size(self)\n",
+        "\n",
         "        for key, val in history.history.items():\n",
         "            aiplatform.log_metrics({f\"train.{key}\": val[-1]})\n",
         "        return history\n",
@@ -793,7 +800,29 @@
         "        aiplatform.log_metrics({\"eval.loss\": metrics[0]})\n",
         "        for _ in range(1, len(metrics)):\n",
         "            aiplatform.log_metrics({\"eval.metric\": metrics[_]})\n",
-        "        return metrics"
+        "        return metrics\n",
+        "\n",
+        "\n",
+        "class TFHelper(object):\n",
+        "    def model_size(self, model):\n",
+        "        \"\"\"\n",
+        "        Get the memory footprint as measured by the number of weights\n",
+        "        \"\"\"\n",
+        "\n",
+        "        def get_size(weights) -> int:\n",
+        "            n = 0\n",
+        "            for weight in weights:\n",
+        "                try:\n",
+        "                    n += len(weight)\n",
+        "                    n += get_size(weight)\n",
+        "                except:\n",
+        "                    pass\n",
+        "\n",
+        "            return n\n",
+        "\n",
+        "        n = get_size(model.get_weights())\n",
+        "        aiplatform.log_metrics({\"n_weights\": n})\n",
+        "        return n"
       ]
     },
     {