Use py_function decorators.

PiperOrigin-RevId: 568972897

Author: MarkDaoust

site/en/guide/data.ipynb

@@ -15,7 +15,6 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "cellView": "form",
         "id": "llMNufAK7nfK"
       },
       "outputs": [],
@@ -589,7 +588,7 @@
       "source": [
         "The first output is an `int32` the second is a `float32`.\n",
         "\n",
-        "The first item is a scalar, shape `()`, and the second is a vector of unknown length, shape `(None,)` "
+        "The first item is a scalar, shape `()`, and the second is a vector of unknown length, shape `(None,)`"
       ]
     },
     {
@@ -601,8 +600,8 @@
       "outputs": [],
       "source": [
         "ds_series = tf.data.Dataset.from_generator(\n",
-        "    gen_series, \n",
-        "    output_types=(tf.int32, tf.float32), \n",
+        "    gen_series,\n",
+        "    output_types=(tf.int32, tf.float32),\n",
         "    output_shapes=((), (None,)))\n",
         "\n",
         "ds_series"
@@ -710,8 +709,8 @@
       "outputs": [],
       "source": [
         "ds = tf.data.Dataset.from_generator(\n",
-        "    lambda: img_gen.flow_from_directory(flowers), \n",
-        "    output_types=(tf.float32, tf.float32), \n",
+        "    lambda: img_gen.flow_from_directory(flowers),\n",
+        "    output_types=(tf.float32, tf.float32),\n",
         "    output_shapes=([32,256,256,3], [32,5])\n",
         ")\n",
         "\n",
@@ -1932,6 +1931,7 @@
       "source": [
         "import scipy.ndimage as ndimage\n",
         "\n",
+        "@tf.py_function(Tout=tf.float32)\n",
         "def random_rotate_image(image):\n",
         "  image = ndimage.rotate(image, np.random.uniform(-30, 30), reshape=False)\n",
         "  return image"
@@ -1969,7 +1969,7 @@
       "source": [
         "def tf_random_rotate_image(image, label):\n",
         "  im_shape = image.shape\n",
-        "  [image,] = tf.py_function(random_rotate_image, [image], [tf.float32])\n",
+        "  image = random_rotate_image(image)\n",
         "  image.set_shape(im_shape)\n",
         "  return image, label"
       ]
@@ -2819,7 +2819,7 @@
         "])\n",
         "\n",
         "model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), \n",
+        "              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])"
       ]
     },
@@ -2953,8 +2953,8 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [],
       "name": "data.ipynb",
+      "provenance": [],
       "toc_visible": true
     },
     "kernelspec": {

site/en/guide/function.ipynb

@@ -1069,7 +1069,66 @@
         "id": "e1I0dPiqTV8H"
       },
       "source": [
-        "If you would like to execute Python code during each invocation of a `Function`, `tf.py_function` is an exit hatch. The drawback of `tf.py_function` is that it's not portable or particularly performant, cannot be saved with SavedModel, and does not work well in distributed (multi-GPU, TPU) setups. Also, since `tf.py_function` has to be wired into the graph, it casts all inputs/outputs to tensors."
+        "If you would like to execute Python code during each invocation of a `Function`, `tf. py_function` is an exit hatch. The drawbacks of `tf.py_function` are that it's not portable or particularly performant, cannot be saved with SavedModel, and does not work well in distributed (multi-GPU, TPU) setups. Also, since `tf.py_function` has to be wired into the graph, it casts all inputs/outputs to tensors."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ZbI7XA_e6yA2"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.py_function(Tout=tf.float32)\n",
+        "def py_plus(x, y):\n",
+        "  print('Executing eagerly.')\n",
+        "  return x + y\n",
+        "\n",
+        "@tf.function\n",
+        "def tf_wrapper(x, y):\n",
+        "  print('Tracing.')\n",
+        "  return py_plus(x, y)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h5ttN_sI7TdQ"
+      },
+      "source": [
+        "The `tf.function` will trace the first time:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "mAK4XINl7Ldy"
+      },
+      "outputs": [],
+      "source": [
+        "tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Atxvrd_o7dSy"
+      },
+      "source": [
+        "But the `tf.py_function` inside executes eagerly every time:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vv7qTiTU7bjy"
+      },
+      "outputs": [],
+      "source": [
+        "tf_wrapper(tf.constant(1.0), tf.constant(2.0)).numpy()"
       ]
     },
     {
@@ -1778,7 +1837,46 @@
         "id": "7Q8BRPCThTjB"
       },
       "source": [
-        "If you need to change the optimizer during training, a workaround is to create a new `Function` for each optimizer, calling the [`ConcreteFunction`](#obtaining_concrete_functions) directly."
+        "If you need to change a stateful object between calls, it's simplest to define a `tf.Module` subclass, and create instances to hold those objects:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3P59ocmIslHz"
+      },
+      "outputs": [],
+      "source": [
+        "class TrainStep(tf.Module):\n",
+        "  def __init__(self, optimizer):\n",
+        "    self.optimizer = optimizer\n",
+        "\n",
+        "  @tf.function\n",
+        "  def __call__(self, w, x, y):\n",
+        "    with tf.GradientTape() as tape:\n",
+        "       L = tf.reduce_sum(tf.square(w*x - y))\n",
+        "    gradients = tape.gradient(L, [w])\n",
+        "    self.optimizer.apply_gradients(zip(gradients, [w]))\n",
+        "\n",
+        "\n",
+        "opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)\n",
+        "opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)\n",
+        "\n",
+        "train_o1 = TrainStep(opt1)\n",
+        "train_o2 = TrainStep(opt2)\n",
+        "\n",
+        "train_o1(w, x, y)\n",
+        "train_o2(w, x, y)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dUHUi881smHF"
+      },
+      "source": [
+        "You could also do this manually by creating multiple instances of the `@tf.function` wrapper, one for each optimizer:"
       ]
     },
     {
@@ -1841,6 +1939,8 @@
   "metadata": {
     "colab": {
       "name": "function.ipynb",
+      "private_outputs": true,
+      "provenance": [],
       "toc_visible": true
     },
     "kernelspec": {