Move the loop out of the tf.function, so this doesn't OOM.

+ collapse code for plotting functions

PiperOrigin-RevId: 420295363

Author: MarkDaoust

site/en/tutorials/interpretability/integrated_gradients.ipynb

@@ -334,6 +334,7 @@
       },
       "outputs": [],
       "source": [
+        "#@title\n",
         "fig = plt.figure(figsize=(12, 5))\n",
         "ax0 = fig.add_subplot(121)\n",
         "ax0.plot(x, f(x), marker='o')\n",
@@ -695,13 +696,23 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "mCH8sAf3TTJ2"
+        "id": "FQWwcI0Wr0AX"
       },
       "outputs": [],
       "source": [
         "pred = model(interpolated_images)\n",
-        "pred_proba = tf.nn.softmax(pred, axis=-1)[:, 555]\n",
-        "\n",
+        "pred_proba = tf.nn.softmax(pred, axis=-1)[:, 555]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "mCH8sAf3TTJ2"
+      },
+      "outputs": [],
+      "source": [
+        "#@title\n",
         "plt.figure(figsize=(10, 4))\n",
         "ax1 = plt.subplot(1, 2, 1)\n",
         "ax1.plot(alphas, pred_proba)\n",
@@ -873,48 +884,59 @@
       },
       "outputs": [],
       "source": [
-        "@tf.function\n",
         "def integrated_gradients(baseline,\n",
         "                         image,\n",
         "                         target_class_idx,\n",
         "                         m_steps=50,\n",
         "                         batch_size=32):\n",
-        "  # 1. Generate alphas.\n",
+        "  # Generate alphas.\n",
         "  alphas = tf.linspace(start=0.0, stop=1.0, num=m_steps+1)\n",
         "\n",
-        "  # Initialize TensorArray outside loop to collect gradients.    \n",
-        "  gradient_batches = tf.TensorArray(tf.float32, size=m_steps+1)\n",
+        "  # Collect gradients.    \n",
+        "  gradient_batches = []\n",
         "    \n",
         "  # Iterate alphas range and batch computation for speed, memory efficiency, and scaling to larger m_steps.\n",
         "  for alpha in tf.range(0, len(alphas), batch_size):\n",
         "    from_ = alpha\n",
         "    to = tf.minimum(from_ + batch_size, len(alphas))\n",
         "    alpha_batch = alphas[from_:to]\n",
         "\n",
-        "    # 2. Generate interpolated inputs between baseline and input.\n",
-        "    interpolated_path_input_batch = interpolate_images(baseline=baseline,\n",
-        "                                                       image=image,\n",
-        "                                                       alphas=alpha_batch)\n",
-        "\n",
-        "    # 3. Compute gradients between model outputs and interpolated inputs.\n",
-        "    gradient_batch = compute_gradients(images=interpolated_path_input_batch,\n",
-        "                                       target_class_idx=target_class_idx)\n",
-        "    \n",
-        "    # Write batch indices and gradients to extend TensorArray.\n",
-        "    gradient_batches = gradient_batches.scatter(tf.range(from_, to), gradient_batch)    \n",
-        "  \n",
+        "    gradient_batch = one_batch(baseline, image, alpha_batch, target_class_idx)\n",
+        "    gradient_batches.append(gradient_batch)\n",
+        "      \n",
         "  # Stack path gradients together row-wise into single tensor.\n",
-        "  total_gradients = gradient_batches.stack()\n",
+        "  total_gradients = tf.stack(gradient_batch)\n",
         "\n",
-        "  # 4. Integral approximation through averaging gradients.\n",
+        "  # Integral approximation through averaging gradients.\n",
         "  avg_gradients = integral_approximation(gradients=total_gradients)\n",
         "\n",
-        "  # 5. Scale integrated gradients with respect to input.\n",
+        "  # Scale integrated gradients with respect to input.\n",
         "  integrated_gradients = (image - baseline) * avg_gradients\n",
         "\n",
         "  return integrated_gradients"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dszwB_Sp0CX0"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def one_batch(baseline, image, alpha_batch, target_class_idx):\n",
+        "    # Generate interpolated inputs between baseline and input.\n",
+        "    interpolated_path_input_batch = interpolate_images(baseline=baseline,\n",
+        "                                                       image=image,\n",
+        "                                                       alphas=alpha_batch)\n",
+        "\n",
+        "    # Compute gradients between model outputs and interpolated inputs.\n",
+        "    gradient_batch = compute_gradients(images=interpolated_path_input_batch,\n",
+        "                                       target_class_idx=target_class_idx)\n",
+        "    return gradient_batch"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -984,6 +1006,7 @@
       },
       "outputs": [],
       "source": [
+        "#@title\n",
         "def plot_img_attributions(baseline,\n",
         "                          image,\n",
         "                          target_class_idx,\n",
@@ -1128,6 +1151,7 @@
     "colab": {
       "collapsed_sections": [],
       "name": "integrated_gradients.ipynb",
+      "private_outputs": true,
       "toc_visible": true
     },
     "kernelspec": {