Update CTL tutorial and guides to make loss reduction consistent

with Keras Model.fit:

  * Divide the prediction loss by actual per-replica batch size
    and num_replicas_in_sync.
  * Always include regularization losses.

This change is for the tensorflow repo. The material on keras.io
can stay as-is.

PiperOrigin-RevId: 558161949

Author: arnoegw

site/en/guide/distributed_training.ipynb

@@ -526,7 +526,9 @@
         "mirrored_strategy = tf.distribute.MirroredStrategy()\n",
         "\n",
         "with mirrored_strategy.scope():\n",
-        "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
+        "  model = tf.keras.Sequential([\n",
+        "      tf.keras.layers.Dense(1, input_shape=(1,),\n",
+        "                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])\n",
         "  model.compile(loss='mse', optimizer='sgd')"
       ]
     },
@@ -673,7 +675,9 @@
       "outputs": [],
       "source": [
         "with mirrored_strategy.scope():\n",
-        "  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])\n",
+        "  model = tf.keras.Sequential([\n",
+        "      tf.keras.layers.Dense(1, input_shape=(1,),\n",
+        "                            kernel_regularizer=tf.keras.regularizers.L2(1e-4))])\n",
         "  optimizer = tf.keras.optimizers.SGD()"
       ]
     },
@@ -716,20 +720,21 @@
       },
       "outputs": [],
       "source": [
+        "# Sets `reduction=NONE` to leave it to tf.nn.compute_average_loss() below.\n",
         "loss_object = tf.keras.losses.BinaryCrossentropy(\n",
         "  from_logits=True,\n",
         "  reduction=tf.keras.losses.Reduction.NONE)\n",
         "\n",
-        "def compute_loss(labels, predictions):\n",
-        "  per_example_loss = loss_object(labels, predictions)\n",
-        "  return tf.nn.compute_average_loss(per_example_loss, global_batch_size=global_batch_size)\n",
-        "\n",
         "def train_step(inputs):\n",
         "  features, labels = inputs\n",
         "\n",
         "  with tf.GradientTape() as tape:\n",
         "    predictions = model(features, training=True)\n",
-        "    loss = compute_loss(labels, predictions)\n",
+        "    per_example_loss = loss_object(labels, predictions)\n",
+        "    loss = tf.nn.compute_average_loss(per_example_loss)\n",
+        "    model_losses = model.losses\n",
+        "    if model_losses:\n",
+        "      loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n",
         "\n",
         "  gradients = tape.gradient(loss, model.trainable_variables)\n",
         "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
@@ -750,9 +755,16 @@
       "source": [
         "A few other things to note in the code above:\n",
         "\n",
-        "1. You used `tf.nn.compute_average_loss` to compute the loss. `tf.nn.compute_average_loss` sums the per example loss and divides the sum by the `global_batch_size`. This is important because later after the gradients are calculated on each replica, they are aggregated across the replicas by **summing** them.\n",
-        "2. You also used the `tf.distribute.Strategy.reduce` API to aggregate the results returned by `tf.distribute.Strategy.run`. `tf.distribute.Strategy.run` returns results from each local replica in the strategy, and there are multiple ways to consume this result. You can `reduce` them to get an aggregated value. You can also do `tf.distribute.Strategy.experimental_local_results` to get the list of values contained in the result, one per local replica.\n",
-        "3. When you call `apply_gradients` within a distribution strategy scope, its behavior is modified. Specifically, before applying gradients on each parallel instance during synchronous training, it performs a sum-over-all-replicas of the gradients.\n"
+        "  1. You used `tf.nn.compute_average_loss` to reduce the per-example prediction losses to a scalar. `tf.nn.compute_average_loss` sums the per example loss and divides the sum by the global batch size. This is important because later after the gradients are calculated on each replica, they are aggregated across the replicas by **summing** them.\n",
+        "\n",
+        "  By default, the global batch size is taken to be `tf.get_strategy().num_replicas_in_sync * tf.shape(per_example_loss)[0]`. It can also be specified explicitly as a keyword argument `global_batch_size=`. Without short batches, the default is equivalent to `tf.nn.compute_average_loss(..., global_batch_size=global_batch_size)` with the `global_batch_size` defined above. (For more on short batches and how to avoid or handle them, see the [Custom Training tutorial](../tutorials/distribute/custom_training.ipynb).)\n",
+        "\n",
+        "  2. You used `tf.nn.scale_regularization_loss` to scale regularization losses registered with the `Model` object, if any, by `1/num_replicas_in_sync` as well. For those regularization losses that are input-dependent, it falls on the modeling code, not the custom training loop, to perform the averaging over the per-replica(!) batch size; that way the modeling code can remain agnostic of replication while the training loop remains agnostic of how regularization losses are computed.\n",
+        "\n",
+        "  3. When you call `apply_gradients` within a distribution strategy scope, its behavior is modified. Specifically, before applying gradients on each parallel instance during synchronous training, it performs a sum-over-all-replicas of the gradients.\n",
+        "\n",
+        "  4. You also used the `tf.distribute.Strategy.reduce` API to aggregate the results returned by `tf.distribute.Strategy.run` for reporting. `tf.distribute.Strategy.run` returns results from each local replica in the strategy, and there are multiple ways to consume this result. You can `reduce` them to get an aggregated value. You can also do `tf.distribute.Strategy.experimental_local_results` to get the list of values contained in the result, one per local replica.\n",
+        "\n"
       ]
     },
     {

site/en/guide/tpu.ipynb

@@ -254,13 +254,32 @@
       "outputs": [],
       "source": [
         "def create_model():\n",
+        "  regularizer = tf.keras.regularizers.L2(1e-5)\n",
         "  return tf.keras.Sequential(\n",
-        "      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),\n",
-        "       tf.keras.layers.Conv2D(256, 3, activation='relu'),\n",
+        "      [tf.keras.layers.Conv2D(256, 3, input_shape=(28, 28, 1),\n",
+        "                              activation='relu',\n",
+        "                              kernel_regularizer=regularizer),\n",
+        "       tf.keras.layers.Conv2D(256, 3,\n",
+        "                              activation='relu',\n",
+        "                              kernel_regularizer=regularizer),\n",
         "       tf.keras.layers.Flatten(),\n",
-        "       tf.keras.layers.Dense(256, activation='relu'),\n",
-        "       tf.keras.layers.Dense(128, activation='relu'),\n",
-        "       tf.keras.layers.Dense(10)])"
+        "       tf.keras.layers.Dense(256,\n",
+        "                             activation='relu',\n",
+        "                             kernel_regularizer=regularizer),\n",
+        "       tf.keras.layers.Dense(128,\n",
+        "                             activation='relu',\n",
+        "                             kernel_regularizer=regularizer),\n",
+        "       tf.keras.layers.Dense(10,\n",
+        "                             kernel_regularizer=regularizer)])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "h-2qaXgfyONQ"
+      },
+      "source": [
+        "This model puts L2 regularization terms on the weights of each layer, so that the custom training loop below can show how you pick them up from `Model.losses`."
       ]
     },
     {
@@ -442,9 +461,13 @@
         "    images, labels = inputs\n",
         "    with tf.GradientTape() as tape:\n",
         "      logits = model(images, training=True)\n",
-        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
+        "      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
         "          labels, logits, from_logits=True)\n",
-        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
+        "      loss = tf.nn.compute_average_loss(per_example_loss)\n",
+        "      model_losses = model.losses\n",
+        "      if model_losses:\n",
+        "        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n",
+        "\n",
         "    grads = tape.gradient(loss, model.trainable_variables)\n",
         "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
         "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
@@ -478,7 +501,7 @@
         "\n",
         "  for step in range(steps_per_epoch):\n",
         "    train_step(train_iterator)\n",
-        "  print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
+        "  print('Current step: {}, training loss: {}, training accuracy: {}%'.format(\n",
         "      optimizer.iterations.numpy(),\n",
         "      round(float(training_loss.result()), 4),\n",
         "      round(float(training_accuracy.result()) * 100, 2)))\n",
@@ -516,9 +539,12 @@
         "    images, labels = inputs\n",
         "    with tf.GradientTape() as tape:\n",
         "      logits = model(images, training=True)\n",
-        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
+        "      per_example_loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
         "          labels, logits, from_logits=True)\n",
-        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
+        "      loss = tf.nn.compute_average_loss(per_example_loss)\n",
+        "      model_losses = model.losses\n",
+        "      if model_losses:\n",
+        "        loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n",
         "    grads = tape.gradient(loss, model.trainable_variables)\n",
         "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
         "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
@@ -531,7 +557,7 @@
         "# retraced if the value changes.\n",
         "train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))\n",
         "\n",
-        "print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
+        "print('Current step: {}, training loss: {}, training accuracy: {}%'.format(\n",
         "      optimizer.iterations.numpy(),\n",
         "      round(float(training_loss.result()), 4),\n",
         "      round(float(training_accuracy.result()) * 100, 2)))"