Add an example of artifact store

.github/workflows/basic_and_faq_usages.yml

@@ -38,5 +38,6 @@ jobs:
         python basic_and_faq_usages/simple_pruning.py
         python basic_and_faq_usages/quadratic_simple_constraint.py
         python basic_and_faq_usages/quadratic_simple_multi_objective.py
+        python basic_and_faq_usages/simple_artifact_store.py
       env:
         OMP_NUM_THREADS: 1

basic_and_faq_usages/requirements.txt

@@ -1,3 +1,4 @@
 optuna
 scikit-learn>=0.19.0
 pandas
+matplotlib

basic_and_faq_usages/simple_artifact_store.py

@@ -0,0 +1,88 @@
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import optuna
+from optuna.artifacts import download_artifact
+from optuna.artifacts import FileSystemArtifactStore
+from optuna.artifacts import get_all_artifact_meta
+from optuna.artifacts import upload_artifact
+import pandas as pd
+
+
+dataset_path = "demo-dataset.csv"
+fig_name = "result-trial.png"
+# NOTE: The uploaded artifacts can be viewed in Optuna Dashboard with the following command:
+# $ optuna-dashboard sqlite:///artifact-demo.db --artifact-dir ./save-artifact-here
+base_path = os.path.join("./save-artifact-here")
+# Make the directory used for artifact store.
+os.makedirs(base_path, exist_ok=True)
+# Instantiate an artifact store.
+artifact_store = FileSystemArtifactStore(base_path=base_path)
+
+
+def create_dataset(dataset_path):
+    # The coefficients we would like to find by Optuna.
+    a_true, b_true = 2, -3
+    X = np.random.random(20)
+    Y = a_true * X + b_true
+    dataset = pd.DataFrame({"X": X, "Y": Y})
+    dataset.to_csv(dataset_path)
+    return dataset
+
+
+dataset = create_dataset(dataset_path)
+
+
+def plot_predictions(a, b, trial):
+    # Create an artifact, which is figure in this example, to upload.
+    _, ax = plt.subplots()
+    x = np.linspace(0, 1, 100)
+    ax.scatter(dataset["X"], dataset["Y"], label="Dataset", color="blue")
+    ax.plot(x, a * x + b, label="Prediction", color="darkred")
+    ax.set_title(f"a={a:.2f}, b={b:.2f}")
+    ax.grid()
+    ax.legend()
+    plt.savefig(fig_name)
+    plt.close()
+
+
+def objective(trial):
+    a = trial.suggest_float("a", -5, 5)
+    b = trial.suggest_float("b", -5, 5)
+    plot_predictions(a, b, trial)
+
+    # Link the plotted figure with trial using artifact store API.
+    upload_artifact(artifact_store=artifact_store, file_path=fig_name, study_or_trial=trial)
+
+    return np.mean((a * dataset["X"] + b - dataset["Y"]) ** 2)
+
+
+def show_best_result(study, artifact_store):
+    best_trial = study.best_trial
+    # Get all the artifact information linked to best_trial. (Here we have only one.)
+    artifact_meta = get_all_artifact_meta(best_trial, storage=study._storage)
+    fig_path = "./result-best-trial.png"
+    # Download the figure from the artifact store to fig_path.
+    download_artifact(
+        artifact_store=artifact_store,
+        artifact_id=artifact_meta[0].artifact_id,
+        file_path=fig_path,
+    )
+    # Display the figure for the best result.
+    best_result_img = plt.imread(fig_path)
+    plt.figure()
+    plt.tick_params(labelleft=False, labelbottom=False, left=False, bottom=False)
+    plt.imshow(best_result_img)
+    plt.show()
+
+
+if __name__ == "__main__":
+    # Create a study with a SQLite storage.
+    study = optuna.create_study(
+        study_name="demo", storage="sqlite:///simple-artifact-store-demo.db", load_if_exists=True
+    )
+    # Upload the dataset to use by artifact store API.
+    upload_artifact(artifact_store=artifact_store, file_path=dataset_path, study_or_trial=study)
+    study.optimize(objective, n_trials=30)
+    show_best_result(study, artifact_store)