|
306 | 306 | "source": [
|
307 | 307 | "*How to do this in TensorFlow?*\n",
|
308 | 308 | "\n",
|
| 309 | + "* Loss reduction and scaling is done automatically in Keras `Model.compile` and `Model.fit`\n", |
| 310 | + "\n", |
309 | 311 | "* If you're writing a custom training loop, as in this tutorial, you should sum the per example losses and divide the sum by the `GLOBAL_BATCH_SIZE`: \n",
|
310 | 312 | "`scale_loss = tf.reduce_sum(loss) * (1. / GLOBAL_BATCH_SIZE)`\n",
|
311 | 313 | "or you can use `tf.nn.compute_average_loss` which takes the per example loss,\n",
|
312 | 314 | "optional sample weights, and `GLOBAL_BATCH_SIZE` as arguments and returns the scaled loss.\n",
|
313 | 315 | "\n",
|
314 |
| - "* If you are using regularization losses in your model then you need to scale\n", |
315 |
| - "the loss value by the number of replicas. You can do this by using the `tf.nn.scale_regularization_loss` function.\n", |
316 |
| - "\n", |
317 |
| - "* Using `tf.reduce_mean` is not recommended. Doing so divides the loss by actual per replica batch size which may vary step to step.\n", |
| 316 | + "* If you're writing a custom training loop for a model with a non-empty list of `Model.losses` (e.g., weight regularizers), you should sum them up and divide the sum by the number of replicas. You can do this by using the `tf.nn.scale_regularization_loss` function.\n", |
318 | 317 | "\n",
|
319 |
| - "* This reduction and scaling is done automatically in Keras `Model.compile` and `Model.fit`\n", |
| 318 | + "* Be careful about batches that are shorter than the `GLOBAL_BATCH_SIZE`, if your training data allows them: Dividing the prediction loss by `GLOBAL_BATCH_SIZE` (instead of using `tf.reduce_mean` over the actual batch size) avoids overweighting examples from short batches. However, this does not apply to regularization losses.\n", |
320 | 319 | "\n",
|
321 | 320 | "* If using `tf.keras.losses` classes (as in the example below), the loss reduction needs to be explicitly specified to be one of `NONE` or `SUM`. `AUTO` and `SUM_OVER_BATCH_SIZE` are disallowed when used with `tf.distribute.Strategy`. `AUTO` is disallowed because the user should explicitly think about what reduction they want to make sure it is correct in the distributed case. `SUM_OVER_BATCH_SIZE` is disallowed because currently it would only divide by per replica batch size, and leave the dividing by number of replicas to the user, which might be easy to miss. So, instead, you need to do the reduction yourself explicitly.\n",
|
322 | 321 | "* If `labels` is multi-dimensional, then average the `per_example_loss` across the number of elements in each sample. For example, if the shape of `predictions` is `(batch_size, H, W, n_classes)` and `labels` is `(batch_size, H, W)`, you will need to update `per_example_loss` like: `per_example_loss /= tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)`\n",
|
|
347 | 346 | " loss_object = tf.keras.losses.SparseCategoricalCrossentropy(\n",
|
348 | 347 | " from_logits=True,\n",
|
349 | 348 | " reduction=tf.keras.losses.Reduction.NONE)\n",
|
350 |
| - " def compute_loss(labels, predictions):\n", |
| 349 | + " def compute_loss(labels, predictions, model_losses):\n", |
351 | 350 | " per_example_loss = loss_object(labels, predictions)\n",
|
352 |
| - " return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)" |
| 351 | + " loss = tf.nn.compute_average_loss(per_example_loss,\n", |
| 352 | + " global_batch_size=GLOBAL_BATCH_SIZE)\n", |
| 353 | + " if model_losses:\n", |
| 354 | + " loss += tf.nn.scale_regularization_loss(tf.add_n(model_losses))\n", |
| 355 | + " return loss" |
353 | 356 | ]
|
354 | 357 | },
|
355 | 358 | {
|
|
419 | 422 | "\n",
|
420 | 423 | " with tf.GradientTape() as tape:\n",
|
421 | 424 | " predictions = model(images, training=True)\n",
|
422 |
| - " loss = compute_loss(labels, predictions)\n", |
| 425 | + " loss = compute_loss(labels, predictions, model.losses)\n", |
423 | 426 | "\n",
|
424 | 427 | " gradients = tape.gradient(loss, model.trainable_variables)\n",
|
425 | 428 | " optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
|
|
0 commit comments