Allowing explicit deferral of setting replication_factor in IPUDataHandler

Summary:
Support passing replication_factor=None into the IPUDataHandler constructor to explicitly defer setting replication_factor.
Previously the best you could do was leave it as the default 1 but this sometime lead to errors validating the values.

TF2.5 Only

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, harrym, christiana

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, christiana

Subscribers: harrym

Maniphest Tasks: T58709

Differential Revision: https://phabricator.sourcevertex.net/D63322

Author: Alfie-Edwards

tensorflow/python/ipu/keras/extensions/data_adapter.py

@@ -48,7 +48,9 @@ def _call_counter(token):  #pylint: disable=missing-type-doc,missing-return-type
 
 
 class IPUDataHandler(data_adapter.DataHandler):
-  """Handles iterating over epoch-level Iterator objects on IPU."""
+  """Handles iterating over epoch-level Iterator objects on IPU.
+  To defer setting the replication factor pass replication_factor=None,
+  then set it later using set_replication_factor."""
   def __init__(  # pylint: disable=super-init-not-called
       self,
       x,
@@ -121,22 +123,27 @@ def __init__(  # pylint: disable=super-init-not-called
           "compatible with the IPU's requirement to know the data shape ahead "
           "of time. Please {}".format(type(x).__name__, hint))
 
-    dataset = self._get_and_post_process_dataset(class_weight)
-
-    self._replication_factor = replication_factor
-    self._inferred_steps = self._infer_steps(steps_per_epoch, dataset)
+    self._dataset = self._get_and_post_process_dataset(class_weight)
     self._steps_per_epoch = steps_per_epoch
+    self._replication_factor = None
+    self._inferred_steps = None
 
-    self._validate_dataset(dataset)
-
-    self._dataset = dataset
     self._current_step = 0
     self._step_increment = self._steps_per_execution_value - 1
     self._insufficient_data = False
 
-    self._validate_data_handler()
+    if replication_factor is not None:
+      self.set_replication_factor(replication_factor)
+
+  def _check_replication_factor_set(self, function_name):
+    if self._replication_factor is None or self._infer_steps is None:
+      raise RuntimeError(
+          f"Cannot call {function_name} before the replication factor is set. "
+          f"Either specify the replication factor in the constructor or set it "
+          f"using set_replication_factor.")
 
   def _validate_data_handler(self):
+    self._check_replication_factor_set("_validate_data_handler")
     super()._validate_data_handler()
 
     if self.steps_per_execution_value > self.inferred_steps:
@@ -212,6 +219,7 @@ def _get_and_post_process_dataset(self, class_weight):
     return dataset
 
   def _validate_dataset(self, dataset):
+    self._check_replication_factor_set("_validate_dataset")
     # Validate the size of the dataset.
     dataset_size = cardinality.cardinality(dataset)
     if dataset_size == cardinality.UNKNOWN:
@@ -263,6 +271,12 @@ def batch_size(self):
     return batch_size
 
   def set_replication_factor(self, value):
+    """Set the replication factor and calculate inferred steps based on the
+    replication factor.
+
+    Args:
+      value (int): The value for the replication factor.
+    """
     self._replication_factor = value
     self._inferred_steps = self._infer_steps(self._steps_per_epoch,
                                              self._dataset)
@@ -271,6 +285,7 @@ def set_replication_factor(self, value):
 
   def enumerate_epochs_with_reuse(self, manager, mode, infeed_kwargs):
     """Yields `(epoch, InfeedQueue)`."""
+    self._check_replication_factor_set("enumerate_epochs_with_reuse")
     with self._truncate_execution_to_epoch():
       data_iterator = manager.get_infeed(mode, self._dataset, infeed_kwargs)
       for epoch in range(self._initial_epoch, self._epochs):

tensorflow/python/ipu/keras/extensions/keras_extension_base.py

@@ -1214,7 +1214,8 @@ def _fit_delegate(self,
           workers=workers,
           use_multiprocessing=use_multiprocessing,
           model=self,
-          steps_per_execution=self._steps_per_execution)
+          steps_per_execution=self._steps_per_execution,
+          replication_factor=None)
 
       # Build the model with specific dtypes. This is important for models
       # without explicit input dtypes (model subclasses and some sequential
@@ -1440,7 +1441,8 @@ def _evaluate_delegate(self,
             workers=workers,
             use_multiprocessing=use_multiprocessing,
             model=self,
-            steps_per_execution=self._steps_per_execution)
+            steps_per_execution=self._steps_per_execution,
+            replication_factor=None)
 
       # Build the model with specific dtypes. This is important for models
       # without explicit input dtypes (model subclasses and some sequential
@@ -1597,7 +1599,8 @@ def _predict_delegate(self,
           workers=workers,
           use_multiprocessing=use_multiprocessing,
           model=self,
-          steps_per_execution=self._steps_per_execution)
+          steps_per_execution=self._steps_per_execution,
+          replication_factor=None)
 
       # Build the model with specific dtypes. This is important for models
       # without explicit input dtypes (model subclasses and some sequential

tensorflow/python/ipu/tests/keras/extensions/data_adapter_test.py

@@ -287,6 +287,19 @@ def test_steps_per_execution_not_divisor(self):
           x, epochs=1, batch_size=2, steps_per_execution=variables.Variable(3))
       del data_handler
 
+  def test_deferred_setting_of_replication_factor(self):
+    x = np.ones((8, 1))
+    x = dataset_ops.Dataset.from_tensor_slices(x).batch(2, drop_remainder=True)
+    # With a batch size of 2, the steps_per_execution value of 3 is not valid
+    # until we set a replication_factor which is a muliple of 3.
+    data_handler = data_adapter.IPUDataHandler(
+        x,
+        epochs=1,
+        batch_size=2,
+        steps_per_execution=variables.Variable(3),
+        replication_factor=None)
+    data_handler.set_replication_factor(3)
+
   def test_dataset_drop_batch_with_replication(self):
     x = np.ones((6, 1))
     x = dataset_ops.Dataset.from_tensor_slices(x).repeat()

tensorflow/python/ipu/tests/keras/keras_poprun_test.py

@@ -32,7 +32,7 @@
 from tensorflow.python.keras import Model, Input
 from tensorflow.python.keras.engine import data_adapter
 from tensorflow.python.keras.optimizer_v2 import gradient_descent
-from tensorflow.python.keras.callbacks import History
+from tensorflow.python.keras.callbacks import Callback, History
 from tensorflow.python.types.core import Tensor
 
 
@@ -210,39 +210,44 @@ def predict_on_full_dataset_without_fixed_size_with_fixed_steps(model):
 
 
 TESTCASES = [
+    # The first arg is the name to use for the test case.
     (
+        "fit_on_full_dataset_with_fixed_size_one_epoch",
         fit_on_full_dataset_with_fixed_size_one_epoch,
         True,
     ),
     (
+        "fit_on_full_dataset_with_fixed_size_two_epochs",
         fit_on_full_dataset_with_fixed_size_two_epochs,
         True,
     ),
     (
+        "evaluate_on_full_dataset_with_fixed_size",
         evaluate_on_full_dataset_with_fixed_size,
         False,
     ),
     (
+        "evaluate_on_full_dataset_with_fixed_size_with_fixed_steps",
         evaluate_on_full_dataset_with_fixed_size_with_fixed_steps,
         False,
     ),
     (
+        "evaluate_on_full_dataset_without_fixed_size_with_fixed_steps",
         evaluate_on_full_dataset_without_fixed_size_with_fixed_steps,
         False,
     ),
     (
+        "predict_on_full_dataset_with_fixed_size",
         predict_on_full_dataset_with_fixed_size,
         False,
     ),
     (
-        predict_on_full_dataset_with_fixed_size,
-        False,
-    ),
-    (
+        "predict_on_full_dataset_with_fixed_size_with_fixed_steps",
         predict_on_full_dataset_with_fixed_size_with_fixed_steps,
         False,
     ),
     (
+        "predict_on_full_dataset_without_fixed_size_with_fixed_steps",
         predict_on_full_dataset_without_fixed_size_with_fixed_steps,
         False,
     ),
@@ -326,7 +331,7 @@ def assert_result_is_equal(self, result_1, result_2):
     else:
       self.assertAllClose(result_1, result_2)
 
-  @parameterized.parameters(*TESTCASES)
+  @parameterized.named_parameters(*TESTCASES)
   def test_popdist_horovod_are_equal(self, callback, did_weights_change):
     """Tests whether the results of using keras from `keras_extensions_base`
     yields the same results as the upstream version after running a callback.
@@ -345,6 +350,27 @@ def test_popdist_horovod_are_equal(self, callback, did_weights_change):
                                          did_weights_change)
     self.assert_result_is_equal(popdist_result, horovod_result)
 
+  def test_popdist_dataset_truncation(self):
+    class StepCounterCallback(Callback):
+      def __init__(self):
+        super().__init__()
+        self.step_count = 0
+
+      def on_train_batch_begin(self, batch, logs=None):
+        self.step_count += 1
+
+    def fn(model):
+      cb = StepCounterCallback()
+      # 5 batches of size 4.
+      dataset = test_dataset(20, batch_size=4)
+      model.fit(dataset, callbacks=[cb], verbose=False)
+      return cb.step_count
+
+    _, result = run_with_popdist(fn)
+
+    # 5 batches should be truncated to 4 batches (2 instances, 2 per instance).
+    self.assertEqual(result, 2)
+
 
 if __name__ == "__main__":
   test.main()