Provide warm-start options and more configurable checkpointing to TF-GNN Runner.

PiperOrigin-RevId: 487367527

Author: Neslihans

tensorflow_gnn/runner/__init__.py

@@ -89,6 +89,7 @@
 # Training
 KerasTrainer = keras_fit.KerasTrainer
 KerasTrainerOptions = keras_fit.KerasTrainerOptions
+KerasTrainerCheckpointOptions = keras_fit.KerasTrainerCheckpointOptions
 
 # Model templates
 ModelFromInitAndUpdates = model_templates.ModelFromInitAndUpdates

tensorflow_gnn/runner/trainers/keras_fit.py

@@ -35,6 +35,26 @@ class KerasTrainerOptions:
   enable_check_numerics: bool = False
 
 
+@dataclasses.dataclass
+class KerasTrainerCheckpointOptions:
+  """Provides Keras Checkpointing related configuration options.
+
+  Attributes:
+    checkpoint_dir: Directory path to save checkpoint files.
+    best_checkpoint: Filename for the best checkpoint.
+    latest_checkpoint: Filename for the latest checkpoint.
+  """
+  checkpoint_dir: Optional[str] = None
+  best_checkpoint: str = "best"
+  latest_checkpoint: str = "latest"
+
+  def best_checkpoint_filepath(self) -> str:
+    return os.path.join(self.checkpoint_dir, self.best_checkpoint)
+
+  def latest_checkpoint_filepath(self) -> str:
+    return os.path.join(self.checkpoint_dir, self.latest_checkpoint)
+
+
 class KerasTrainer:
   """Trains using the `tf.keras.Model.fit` training loop."""
 
@@ -43,7 +63,7 @@ def __init__(
       strategy: tf.distribute.Strategy,
       *,
       model_dir: str,
-      ckpts_dir: Optional[str] = None,
+      checkpoint_options: Optional[KerasTrainerCheckpointOptions] = None,
       backup_dir: Optional[str] = None,
       steps_per_epoch: Optional[int] = None,
       validation_steps: Optional[int] = None,
@@ -58,8 +78,9 @@ def __init__(
     Args:
       strategy: A `tf.distribute.Strategy.`
       model_dir: A model directory for summaries.
-      ckpts_dir: An optional directory for checkpoints, if unset;
-        `os.path.join(model_dir, "ckpts")` is used.
+      checkpoint_options: An optional configuration for checkpointing related
+        configs. If checkpoint_options.checkpoint_dir is unset;
+        `os.path.join(model_dir, "ckpnt")` is used.
       backup_dir: An optional directory for backup, if unset;
         `(os.path.join(model_dir, "backup"),)` is used.
       steps_per_epoch: An optional steps per epoch, if unspecified: epochs are
@@ -87,15 +108,16 @@ def __init__(
       raise ValueError("`restore_best_weights` requires a "
                        "`checkpoint_every_n_steps` other than \"never\"")
 
-    if ckpts_dir is None:
-      ckpts_dir = os.path.join(model_dir, "ckpts")
+    if checkpoint_options is None:
+      checkpoint_options = KerasTrainerCheckpointOptions()
+      checkpoint_options.checkpoint_dir = os.path.join(model_dir, "ckpnt")
 
     if backup_dir is None:
       backup_dir = os.path.join(model_dir, "backup")
 
     self._strategy = strategy
     self._model_dir = model_dir
-    self._ckpts_dir = ckpts_dir
+    self._checkpoint_options = checkpoint_options
     self._backup_dir = backup_dir
     self._steps_per_epoch = steps_per_epoch
     self._validation_steps = validation_steps
@@ -215,12 +237,12 @@ def per_replica_ds_fn(input_context, *, delegate, repeat):
     if checkpoint_every_n_steps != "never":
       callbacks += [
           tf.keras.callbacks.ModelCheckpoint(
-              filepath=os.path.join(self._ckpts_dir, "latest"),
+              filepath=self._checkpoint_options.latest_checkpoint_filepath(),
               save_best_only=False,
               save_weights_only=True,
               save_freq=checkpoint_every_n_steps),
           tf.keras.callbacks.ModelCheckpoint(
-              filepath=os.path.join(self._ckpts_dir, "best"),
+              filepath=self._checkpoint_options.best_checkpoint_filepath(),
               save_best_only=True,
               save_weights_only=True,
               save_freq="epoch")
@@ -256,7 +278,6 @@ def per_replica_ds_fn(input_context, *, delegate, repeat):
         callbacks=callbacks)
 
     if self._restore_best_weights:
-      model.load_weights(os.path.join(self._ckpts_dir, "best"))
+      model.load_weights(self._checkpoint_options.best_checkpoint_filepath())
 
     return model
-