Change examples

examples/preferential-optimization/evaluator.py

@@ -1,94 +0,0 @@
-from __future__ import annotations
-
-import os
-import shutil
-import tempfile
-import time
-from typing import Callable
-from typing import NoReturn
-import uuid
-
-from optuna_dashboard.artifact.file_system import FileSystemBackend
-from optuna_dashboard.preferential import load_study
-
-import streamlit as st
-
-
-STORAGE_URL = "sqlite:///st-example.db"
-artifact_path = os.path.join(os.path.dirname(__file__), "artifact")
-artifact_backend = FileSystemBackend(base_path=artifact_path)
-os.makedirs(artifact_path, exist_ok=True)
-
-n_comparison = 5
-
-
-def get_tmp_dir() -> str:
-    if "tmp_dir" not in st.session_state:
-        tmp_dir_name = str(uuid.uuid4())
-        tmp_dir_path = os.path.join(tempfile.gettempdir(), tmp_dir_name)
-        os.makedirs(tmp_dir_path, exist_ok=True)
-        st.session_state.tmp_dir = tmp_dir_path
-
-    return st.session_state.tmp_dir
-
-
-def main() -> NoReturn:
-    tmpdir = get_tmp_dir()
-    study = load_study(
-        study_name="Preferential Optimization",
-        storage=STORAGE_URL,
-    )
-
-    # 1. Get all currently best trials (i.e. trials that are not reported bad) for comparison.
-    comparison_trials = study.best_trials
-
-    st.text("Which is the worst?")
-
-    # 2. Show the artifact images of all those trials.
-    cols = st.columns(n_comparison)
-    finished_dict = {t.number: t for t in comparison_trials}
-
-    col_is: dict[int, int] = st.session_state.get("col_is")
-    if col_is is None:
-        col_is = {}
-    col_is = {tn: col_i for (tn, col_i) in col_is.items() if tn in finished_dict}
-
-    unoccupied_col_is = [i for i in range(len(cols)) if i not in col_is.values()]
-    for tn, col_i in zip([tn for tn in finished_dict if tn not in col_is], unoccupied_col_is):
-        col_is[tn] = col_i
-    st.session_state["col_is"] = col_is
-
-    def on_click_factory(trial_number: int) -> Callable[[], None]:
-        def on_click() -> None:
-            better_trials = [t for t in comparison_trials if t.number != trial_number]
-            worse_trial = finished_dict[trial_number]
-            study.report_preference(better_trials, worse_trial)
-
-        return on_click
-
-    for trial_number, col_i in col_is.items():
-        trial = finished_dict[trial_number]
-        col = cols[col_i]
-
-        rgb_artifact_id = trial.user_attrs["rgb_artifact_id"]
-        image_caption = trial.user_attrs["image_caption"]
-        with col:
-            with artifact_backend.open(rgb_artifact_id) as fsrc:
-                tmp_img_path = os.path.join(tmpdir, rgb_artifact_id + ".png")
-                with open(tmp_img_path, "wb") as fdst:
-                    shutil.copyfileobj(fsrc, fdst)
-            st.image(tmp_img_path, caption=image_caption)
-            st.button(str(trial_number), key=trial.number, on_click=on_click_factory(trial_number))
-
-    for i, col in enumerate(st.columns(n_comparison)):
-        if i >= len(comparison_trials):
-            continue
-
-    if len(comparison_trials) < n_comparison:
-        # Wait for unfinished trials (images under generation) to be generated.
-        time.sleep(0.1)
-        st.experimental_rerun()
-
-
-if __name__ == "__main__":
-    main()

examples/preferential-optimization/evaluator.sh

@@ -0,0 +1,2 @@
+#!/usr/bin/env sh
+optuna-dashboard sqlite:///example.db --artifact-dir ./artifact

examples/preferential-optimization/generator.py

@@ -7,14 +7,14 @@ import time
 from typing import NoReturn
 
 from optuna_dashboard import save_note
-from optuna_dashboard.artifact import upload_artifact
+from optuna_dashboard.artifact import upload_artifact, get_artifact_path
 from optuna_dashboard.artifact.file_system import FileSystemBackend
 from optuna_dashboard.preferential import create_study
 from optuna_dashboard.preferential.samplers._gp import PreferentialGPSampler
 from PIL import Image
 
 
-STORAGE_URL = "sqlite:///st-example.db"
+STORAGE_URL = "sqlite:///example.db"
 artifact_path = os.path.join(os.path.dirname(__file__), "artifact")
 artifact_backend = FileSystemBackend(base_path=artifact_path)
 os.makedirs(artifact_path, exist_ok=True)
@@ -51,14 +51,14 @@ def main() -> NoReturn:
 
             # 3. Upload Artifact
             artifact_id = upload_artifact(artifact_backend, trial, image_path)
-            trial.set_user_attr("rgb_artifact_id", artifact_id)
-            trial.set_user_attr("image_caption", f"(R, G, B) = ({r}, {g}, {b})")
+            trial.set_user_attr("artifact_id", artifact_id)
             print("RGB:", (r, g, b))
 
             # 4. Save Note
             note = textwrap.dedent(
                 f"""\
-            ![generated-image]({artifact_path})
+            ![generated-image]({get_artifact_path(trial, artifact_id)})
+            
             (R, G, B) = ({r}, {g}, {b})
             """
             )

optuna_dashboard/preferential/samplers/_gp.py

@@ -31,6 +31,7 @@ import pyro.infer.mcmc
 from scipy.special import erfcinv
 import torch
 from torch import Tensor
+from linear_operator.utils.errors import NotPSDError
 
 from .._system_attrs import get_preferences
 
@@ -181,7 +182,9 @@ class _PreferentialGP(GPyTorchModel, ExactGP):
     def _pyro_model(self, train_x: torch.Tensor, train_y: torch.Tensor) -> None:
         # with gpytorch.settings.fast_computations(False, False, False):
         sampled_model = self.pyro_sample_from_prior()
+        
         ys = sampled_model.likelihood(sampled_model.forward(train_x))
+        
         pyro.sample("y", ys, obs=train_y)
 
     def fit_mcmc(self, X: torch.Tensor, preferences: torch.Tensor, cycles: int, rng: np.random.RandomState) -> None:
@@ -244,7 +247,19 @@ class _PreferentialGP(GPyTorchModel, ExactGP):
                 ys_sum = torch.from_numpy(ys_sum_np)
                 train_y[:] = ys_sum[mask] / cnt[mask]
                 nuts.clear_cache()
-                raw_params = nuts.sample(raw_params)
+                try:
+                    raw_params = nuts.sample(raw_params)
+                except NotPSDError:
+                    nuts.cleanup()
+                    nuts = pyro.infer.mcmc.NUTS(
+                        model=self._pyro_model,
+                        init_strategy=pyro.infer.autoguide.init_to_sample,
+                        step_size=self._last_mcmc_step_size or 1.0,
+                    )
+                    nuts.setup(warmup_steps=warmup_steps, train_x=train_x, train_y=train_y)
+                    raw_params = nuts.initial_params
+
+                    
 
             params = {name: nuts.transforms[name].inv(value) for name, value in raw_params.items()}
             self.set_train_data(train_x, train_y, strict=False)
@@ -326,7 +341,8 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler):
             if len(search_space) == 0:
                 return {}
             
-            preferences = get_preferences(study, deepcopy=False)
+            preferences = get_preferences(study._study_id, study._storage)
+            trials = study.get_trials(deepcopy=False)
             if len(preferences) == 0:
                 return {}
 
@@ -352,10 +368,10 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler):
 
             for better, worse in preferences:
                 for t in (better, worse):
-                    if t.number not in ids:
-                        ids[t.number] = len(ids)
-                        params.append(trans.transform(t.params))
-                pref_ids.append((ids[better.number], ids[worse.number]))
+                    if t not in ids:
+                        ids[t] = len(ids)
+                        params.append(trans.transform(trials[t].params))
+                pref_ids.append((ids[better], ids[worse]))
             dtype = torch.float64
 
             params_torch = torch.tensor(np.array(params), dtype=dtype, device=self.device)