Introduce optuna.artifacts to the PyTorch checkpoint example

[kAIto47802]

Motivation

Currently, the PyTorch checkpoint example is using local file system to save and manage checkpoints, not yet reflecting the recent optuna.artifacts functionalities.

Description of the changes

• Introduced optuna.artifacts.
• Removed the use of local file system.

[HideakiImamura]

@nabenabe0928 Could you review this PR?

[github-actions]

This pull request has not seen any recent activity.

[not522]

Thank you for your PR! Could you check my comments?

[not522]

If the process is terminated before the first checkpoint, the artifact will not be saved, so check if it exists.

[not522]

Could you change the path of checkpoint for each trial? If we run this script with multi-process, the saved models can be broken by other processes.

[github-actions]

This pull request has not seen any recent activity.

[not522]

retried_trial_number returns the first trial's number in the retry history, so we should check the entire retry history. For example, running the following command can check the behavior.

timeout 5 python pytorch_checkpoint.py

The fix could be like this.

$ git diff
diff --git a/pytorch/pytorch_checkpoint.py b/pytorch/pytorch_checkpoint.py
index 35b697d..e7f2b5d 100644
--- a/pytorch/pytorch_checkpoint.py
+++ b/pytorch/pytorch_checkpoint.py
@@ -89,20 +89,25 @@ def objective(trial):
     lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
     optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
 
-    trial_number = RetryFailedTrialCallback.retried_trial_number(trial)
-
-    artifact_id = trial_number and trial.study.trials[trial_number].user_attrs.get("artifact_id")
-    if trial_number is not None and artifact_id is not None:
+    artifact_id = None
+    retry_history = RetryFailedTrialCallback.retry_history(trial)
+    for trial_number in reversed(retry_history):
+        artifact_id = trial.study.trials[trial_number].user_attrs.get("artifact_id")
+        if artifact_id is not None:
+            retry_trial_number = trial_number
+            break
+
+    if artifact_id is not None:
         download_artifact(
             artifact_store=artifact_store,
-            file_path=f"./tmp_model_{trial_number}.pt",
+            file_path=f"./tmp_model_{trial.number}.pt",
             artifact_id=artifact_id,
         )
-        checkpoint = torch.load(f"./tmp_model_{trial_number}.pt")
+        checkpoint = torch.load(f"./tmp_model_{trial.number}.pt")
         epoch = checkpoint["epoch"]
         epoch_begin = epoch + 1
 
-        print(f"Loading a checkpoint from trial {trial_number} in epoch {epoch}.")
+        print(f"Loading a checkpoint from trial {retry_trial_number} in epoch {epoch}.")
 
         model.load_state_dict(checkpoint["model_state_dict"])
         optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
@@ -159,15 +164,15 @@ def objective(trial):
                 "optimizer_state_dict": optimizer.state_dict(),
                 "accuracy": accuracy,
             },
-            f"./tmp_model_{trial_number}.pt",
+            f"./tmp_model_{trial.number}.pt",
         )
         artifact_id = upload_artifact(
             artifact_store=artifact_store,
-            file_path=f"./tmp_model_{trial_number}.pt",
+            file_path=f"./tmp_model_{trial.number}.pt",
             study_or_trial=trial,
         )
         trial.set_user_attr("artifact_id", artifact_id)
-        os.remove(f"./tmp_model_{trial_number}.pt")
+        os.remove(f"./tmp_model_{trial.number}.pt")
 
         # Handle pruning based on the intermediate value.
         if trial.should_prune():

[kAIto47802]

Thank you for your review! I have fixed it according to your suggestion.

[not522]

Thank you for your update. It's almost LGTM. Could you check my comment?

[not522]

Could you remove the temporary file here?

os.remove(f"./tmp_model_{trial.number}.pt")

[kAIto47802]

Thank you for your comment. I have fix this.

[not522]

LGTM!