Add guidance on enabling exactly-once evaluation for ParameterServerStrategy.

PiperOrigin-RevId: 525146158

Author: jamesmullenbach

site/en/tutorials/distribute/parameter_server_training.ipynb

@@ -132,10 +132,9 @@
         "\n",
         "With custom training loops, the `tf.distribute.coordinator.ClusterCoordinator` class is the key component used for the coordinator.\n",
         "\n",
-        "- The `ClusterCoordinator` class needs to work in conjunction with a `tf.distribute.Strategy` object.\n",
+        "- The `ClusterCoordinator` class needs to work in conjunction with a `tf.distribute.ParameterServerStrategy` object.\n",
         "- This `tf.distribute.Strategy` object is needed to provide the information of the cluster and is used to define a training step, as demonstrated in [Custom training with tf.distribute.Strategy](custom_training.ipynb).\n",
         "- The `ClusterCoordinator` object then dispatches the execution of these training steps to remote workers.\n",
-        "- For parameter server training, the `ClusterCoordinator` needs to work with a `tf.distribute.ParameterServerStrategy`.\n",
         "\n",
         "The most important API provided by the `ClusterCoordinator` object is `schedule`:\n",
         "\n",
@@ -883,7 +882,7 @@
       "source": [
         "## Evaluation\n",
         "\n",
-        "The two main approaches to performing evaluation with `tf.distribute.ParameterServerStrategy` training are inline evaluation and sidecar evaluation. Each has its own pros and cons as described below. The inline evaluation method is recommended if you don't have a preference."
+        "The two main approaches to performing evaluation with `tf.distribute.ParameterServerStrategy` training are inline evaluation and sidecar evaluation. Each has its own pros and cons as described below. The inline evaluation method is recommended if you don't have a preference. For users using `Model.fit`, `Model.evaluate` uses inline (distributed) evaluation under the hood."
       ]
     },
     {
@@ -985,7 +984,23 @@
         "id": "cKrQktZX5z7a"
       },
       "source": [
-        "Note: The `schedule` and `join` methods of `tf.distribute.coordinator.ClusterCoordinator` don’t support visitation guarantees or exactly-once semantics. In other words, there is no guarantee that all evaluation examples in a dataset will be evaluated exactly once; some may not be visited and some may be evaluated multiple times. The `tf.data` service API can be used to provide exactly-once visitation for evaluation when using `ParameterServerStrategy` (refer to the _Dynamic Sharding_ section of the `tf.data.experimental.service` API documentation)."
+        "#### Enabling exactly-once evaluation\n",
+        "<a id=\"exactly_once_evaluation\"></a>\n",
+        "\n",
+        "The `schedule` and `join` methods of `tf.distribute.coordinator.ClusterCoordinator` don’t support visitation guarantees or exactly-once semantics by default. In other words, in the above example there is no guarantee that all evaluation examples in a dataset will be evaluated exactly once; some may not be visited and some may be evaluated multiple times. \n",
+        "\n",
+        "Exactly-once evaluation may be preferred to reduce the variance of evaluation across epochs, and improve model selection done via early stopping,  hyperparameter tuning, or other methods. There are different ways to enable exactly-once evaluation:\n",
+        "\n",
+        "- With a `Model.fit/.evaluate` workflow, it can be enabled by adding an argument to `Model.compile`. Refer to docs for the `pss_evaluation_shards` argument.\n",
+        "- The `tf.data` service API can be used to provide exactly-once visitation for evaluation when using `ParameterServerStrategy` (refer to the _Dynamic Sharding_ section of the `tf.data.experimental.service` API documentation).\n",
+        "- [Sidecar evaluation](#sidecar_evaluation) provides exactly-once evaluation by default, since the evaluation happens on a single machine. However this can be much slower than performing evaluation distributed across many workers.\n",
+        "\n",
+        "The first option, using `Model.compile`, is the suggested solution for most users. \n",
+        "\n",
+        "Exactly-once evaluation has some limitations:\n",
+        "\n",
+        "- It is not supported to write a custom distributed evaluation loop with an exactly-once visitation guarantee. File a GitHub issue if you need support for this.\n",
+        "- It cannot automatically handle computation of metrics that use the `Layer.add_metric` API. These should be excluded from evaluation, or reworked into `Metric` objects."
       ]
     },
     {
@@ -997,7 +1012,7 @@
         "### Sidecar evaluation\n",
         "<a id=\"sidecar_evaluation\"></a>\n",
         "\n",
-        "Another method for defining and running an evaluation loop in `tf.distribute.ParameterServerStrategy` training is called _sidecar evaluation_, in which you create a dedicated evaluator task that repeatedly reads checkpoints and runs evaluation on the latest checkpoint (refer to [this guide](../../guide/checkpoint.ipynb) for more details on checkpointing). The chief and worker tasks do not spend any time on evaluation, so for a fixed number of iterations the overall training time should be shorter than using other evaluation methods. However, it requires an additional evaluator task and periodic checkpointing to trigger evaluation."
+        "Another method for defining and running an evaluation loop in `tf.distribute.ParameterServerStrategy` training is called _sidecar evaluation_, in which you create a dedicated evaluator task that repeatedly reads checkpoints and runs evaluation on the latest checkpoint (refer to [this guide](../../guide/checkpoint.ipynb) for more details on checkpointing). The coordinator and worker tasks do not spend any time on evaluation, so for a fixed number of iterations the overall training time should be shorter than using other evaluation methods. However, it requires an additional evaluator task and periodic checkpointing to trigger evaluation."
       ]
     },
     {
@@ -1348,7 +1363,7 @@
       "source": [
         "### Custom training loop specifics\n",
         "\n",
-        "- `ClusterCoordinator.schedule` doesn't support visitation guarantees for a dataset.\n",
+        "- `ClusterCoordinator.schedule` doesn't support visitation guarantees for a dataset in general, although a visitation guarantee for evaluation is possible through `Model.fit/.evaluate`. See [Enabling exactly-once evaluation](#exactly_once_evaluation).\n",
         "- When `ClusterCoordinator.create_per_worker_dataset` is used with a callable as input, the whole dataset must be created inside the function passed to it.\n",
         "- `tf.data.Options` is ignored in a dataset created by `ClusterCoordinator.create_per_worker_dataset`."
       ]
@@ -1357,7 +1372,6 @@
   "metadata": {
     "accelerator": "GPU",
     "colab": {
-      "collapsed_sections": [],
       "name": "parameter_server_training.ipynb",
       "provenance": [],
       "toc_visible": true