[train][doc] New checkpointing user guide (ray-project#39505)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>

Author: justinvyu

doc/source/train/distributed-tensorflow-keras.rst

@@ -230,7 +230,9 @@ appropriately in distributed training.
 
 
 .. code-block:: python
-    :emphasize-lines: 23
+
+    import os
+    import tempfile
 
     from ray import train
     from ray.train import Checkpoint, ScalingConfig
@@ -254,24 +256,24 @@ appropriately in distributed training.
             model.compile(optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
 
         for epoch in range(config["num_epochs"]):
-            model.fit(X, Y, batch_size=20)
-            checkpoint = Checkpoint.from_dict(
-                dict(epoch=epoch, model_weights=model.get_weights())
-            )
-            train.report({}, checkpoint=checkpoint)
+            history = model.fit(X, Y, batch_size=20)
+
+            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                model.save(os.path.join(temp_checkpoint_dir, "model.keras"))
+                checkpoint_dict = os.path.join(temp_checkpoint_dir, "checkpoint.json")
+                with open(checkpoint_dict, "w") as f:
+                    json.dump({"epoch": epoch}, f)
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+                train.report({"loss": history.history["loss"][0]}, checkpoint=checkpoint)
 
     trainer = TensorflowTrainer(
         train_func,
         train_loop_config={"num_epochs": 5},
         scaling_config=ScalingConfig(num_workers=2),
     )
     result = trainer.fit()
-
-    print(result.checkpoint.to_dict())
-    # {'epoch': 4, 'model_weights': [array([[-0.31858477],
-    #    [ 0.03747174],
-    #    [ 0.28266194],
-    #    [ 0.8626015 ]], dtype=float32), array([0.02230084], dtype=float32)], '_timestamp': 1656107383, '_preprocessor': None, '_current_checkpoint_id': 4}
+    print(result.checkpoint)
 
 By default, checkpoints will be persisted to local disk in the :ref:`log
 directory <train-log-dir>` of each run.
@@ -280,7 +282,9 @@ Loading checkpoints
 ~~~~~~~~~~~~~~~~~~~
 
 .. code-block:: python
-    :emphasize-lines: 15, 21, 22, 25, 26, 27, 30
+
+    import os
+    import tempfile
 
     from ray import train
     from ray.train import Checkpoint, ScalingConfig
@@ -297,37 +301,42 @@ Loading checkpoints
         X = np.random.normal(0, 1, size=(n, 4))
         Y = np.random.uniform(0, 1, size=(n, 1))
 
-        start_epoch = 0
         strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
-
         with strategy.scope():
             # toy neural network : 1-layer
-            model = tf.keras.Sequential([tf.keras.layers.Dense(1, activation="linear", input_shape=(4,))])
             checkpoint = train.get_checkpoint()
             if checkpoint:
-                # assume that we have run the train.report() example
-                # and successfully save some model weights
-                checkpoint_dict = checkpoint.to_dict()
-                model.set_weights(checkpoint_dict.get("model_weights"))
-                start_epoch = checkpoint_dict.get("epoch", -1) + 1
+                with checkpoint.as_directory() as checkpoint_dir:
+                    model = tf.keras.models.load_model(
+                        os.path.join(checkpoint_dir, "model.keras")
+                    )
+            else:
+                model = tf.keras.Sequential(
+                    [tf.keras.layers.Dense(1, activation="linear", input_shape=(4,))]
+                )
             model.compile(optimizer="Adam", loss="mean_squared_error", metrics=["mse"])
 
-        for epoch in range(start_epoch, config["num_epochs"]):
-            model.fit(X, Y, batch_size=20)
-            checkpoint = Checkpoint.from_dict(
-                dict(epoch=epoch, model_weights=model.get_weights())
-            )
-            train.report({}, checkpoint=checkpoint)
+        for epoch in range(config["num_epochs"]):
+            history = model.fit(X, Y, batch_size=20)
+
+            with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
+                model.save(os.path.join(temp_checkpoint_dir, "model.keras"))
+                extra_json = os.path.join(temp_checkpoint_dir, "checkpoint.json")
+                with open(extra_json, "w") as f:
+                    json.dump({"epoch": epoch}, f)
+                checkpoint = Checkpoint.from_directory(temp_checkpoint_dir)
+
+                train.report({"loss": history.history["loss"][0]}, checkpoint=checkpoint)
 
     trainer = TensorflowTrainer(
         train_func,
-        train_loop_config={"num_epochs": 2},
+        train_loop_config={"num_epochs": 5},
         scaling_config=ScalingConfig(num_workers=2),
     )
-    # save a checkpoint
     result = trainer.fit()
+    print(result.checkpoint)
 
-    # load a checkpoint
+    # Start a new run from a loaded checkpoint
     trainer = TensorflowTrainer(
         train_func,
         train_loop_config={"num_epochs": 5},
@@ -336,12 +345,6 @@ Loading checkpoints
     )
     result = trainer.fit()
 
-    print(result.checkpoint.to_dict())
-    # {'epoch': 4, 'model_weights': [array([[-0.70056134],
-    #    [-0.8839263 ],
-    #    [-1.0043601 ],
-    #    [-0.61634773]], dtype=float32), array([0.01889327], dtype=float32)], '_timestamp': 1656108446, '_preprocessor': None, '_current_checkpoint_id': 3}
-
 
 Further reading
 ---------------