Fix Checkpoint in Hyperparameter Tuning

beginner_source/hyperparameter_tuning_tutorial.py

@@ -41,6 +41,8 @@
 """
 from functools import partial
 import os
+import tempfile
+from pathlib import Path
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -49,11 +51,14 @@
 import torchvision
 import torchvision.transforms as transforms
 from ray import tune
-from ray.air import Checkpoint, session
+from ray import train
+from ray.train import Checkpoint, get_checkpoint
+from ray.air import Checkpoint
 from ray.tune.schedulers import ASHAScheduler
+import ray.cloudpickle as pickle
 
 ######################################################################
-# Most of the imports are needed for building the PyTorch model. Only the last three
+# Most of the imports are needed for building the PyTorch model. Only the last 
 # imports are for Ray Tune.
 #
 # Data loaders
@@ -124,13 +129,15 @@ def forward(self, x):
 #
 #     net = Net(config["l1"], config["l2"])
 #
-#     checkpoint = session.get_checkpoint()
-#
+#     checkpoint = get_checkpoint()
 #     if checkpoint:
-#         checkpoint_state = checkpoint.to_dict()
-#         start_epoch = checkpoint_state["epoch"]
-#         net.load_state_dict(checkpoint_state["net_state_dict"])
-#         optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"])
+#         with train.get_checkpoint().as_directory() as checkpoint_dir:
+#             data_path = Path(checkpoint_dir) / "data.pkl"
+#             with open(data_path, "rb") as fp:
+#                 checkpoint_state = pickle.load(fp)
+#                 start_epoch = checkpoint_state["epoch"]
+#                 net.load_state_dict(checkpoint_state["net_state_dict"])
+#                 optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"])
 #     else:
 #         start_epoch = 0
 #
@@ -186,12 +193,16 @@ def forward(self, x):
 #         "net_state_dict": net.state_dict(),
 #         "optimizer_state_dict": optimizer.state_dict(),
 #     }
-#     checkpoint = Checkpoint.from_dict(checkpoint_data)
+#     with tempfile.TemporaryDirectory() as checkpoint_dir:
+#         data_path = Path(checkpoint_dir) / "data.pkl"
+#         with open(data_path, "wb") as fp:
+#             pickle.dump(checkpoint_data, fp)
 #
-#     session.report(
-#         {"loss": val_loss / val_steps, "accuracy": correct / total},
-#         checkpoint=checkpoint,
-#     )
+#             checkpoint = Checkpoint.from_directory(checkpoint_dir)
+#             train.report(
+#                 {"loss": val_loss / val_steps, "accuracy": correct / total},
+#                 checkpoint=checkpoint,
+#             )
 #
 # Here we first save a checkpoint and then report some metrics back to Ray Tune. Specifically,
 # we send the validation loss and accuracy back to Ray Tune. Ray Tune can then use these metrics
@@ -225,13 +236,15 @@ def train_cifar(config, data_dir=None):
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(net.parameters(), lr=config["lr"], momentum=0.9)
 
-    checkpoint = session.get_checkpoint()
-
+    checkpoint = get_checkpoint()
     if checkpoint:
-        checkpoint_state = checkpoint.to_dict()
-        start_epoch = checkpoint_state["epoch"]
-        net.load_state_dict(checkpoint_state["net_state_dict"])
-        optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"])
+        with train.get_checkpoint().as_directory() as checkpoint_dir:
+            data_path = Path(checkpoint_dir) / "data.pkl"
+            with open(data_path, "rb") as fp:
+                checkpoint_state = pickle.load(fp)
+                start_epoch = checkpoint_state["epoch"]
+                net.load_state_dict(checkpoint_state["net_state_dict"])
+                optimizer.load_state_dict(checkpoint_state["optimizer_state_dict"])
     else:
         start_epoch = 0
 
@@ -300,12 +313,17 @@ def train_cifar(config, data_dir=None):
             "net_state_dict": net.state_dict(),
             "optimizer_state_dict": optimizer.state_dict(),
         }
-        checkpoint = Checkpoint.from_dict(checkpoint_data)
-
-        session.report(
-            {"loss": val_loss / val_steps, "accuracy": correct / total},
-            checkpoint=checkpoint,
-        )
+        with tempfile.TemporaryDirectory() as checkpoint_dir:
+            data_path = Path(checkpoint_dir) / "data.pkl"
+            with open(data_path, "wb") as fp:
+                pickle.dump(checkpoint_data, fp)
+
+            checkpoint = Checkpoint.from_directory(checkpoint_dir)
+            train.report(
+                {"loss": val_loss / val_steps, "accuracy": correct / total},
+                checkpoint=checkpoint,
+            )
+
     print("Finished Training")
 
 
@@ -438,13 +456,14 @@ def main(num_samples=10, max_num_epochs=10, gpus_per_trial=2):
             best_trained_model = nn.DataParallel(best_trained_model)
     best_trained_model.to(device)
 
-    best_checkpoint = best_trial.checkpoint.to_air_checkpoint()
-    best_checkpoint_data = best_checkpoint.to_dict()
-
-    best_trained_model.load_state_dict(best_checkpoint_data["net_state_dict"])
-
-    test_acc = test_accuracy(best_trained_model, device)
-    print("Best trial test set accuracy: {}".format(test_acc))
+    best_checkpoint = best_trial.checkpoint
+    with best_checkpoint.as_directory() as checkpoint_dir:
+        data_path = Path(checkpoint_dir) / "data.pkl"
+        with open(data_path, "rb") as fp:
+            best_checkpoint_data = pickle.load(fp)
+            best_trained_model.load_state_dict(best_checkpoint_data["net_state_dict"])
+            test_acc = test_accuracy(best_trained_model, device)
+            print("Best trial test set accuracy: {}".format(test_acc))
 
 
 if __name__ == "__main__":