From 744bdee4de35b1408ed8c80ccec39c499c6a5253 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Tue, 22 Aug 2023 19:27:25 +0900 Subject: [PATCH] Change examples --- .../preferential-optimization/evaluator.py | 94 ------------------- .../preferential-optimization/evaluator.sh | 2 + .../preferential-optimization/generator.py | 10 +- optuna_dashboard/preferential/samplers/_gp.py | 28 ++++-- 4 files changed, 29 insertions(+), 105 deletions(-) delete mode 100644 examples/preferential-optimization/evaluator.py create mode 100755 examples/preferential-optimization/evaluator.sh diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py deleted file mode 100644 index ee1950c5..00000000 --- a/examples/preferential-optimization/evaluator.py +++ /dev/null @@ -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() diff --git a/examples/preferential-optimization/evaluator.sh b/examples/preferential-optimization/evaluator.sh new file mode 100755 index 00000000..f7ebeca5 --- /dev/null +++ b/examples/preferential-optimization/evaluator.sh @@ -0,0 +1,2 @@ +#!/usr/bin/env sh +optuna-dashboard sqlite:///example.db --artifact-dir ./artifact \ No newline at end of file diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index 07977d0b..bbfea678 100644 --- a/examples/preferential-optimization/generator.py +++ b/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}) """ ) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index 39d85b5b..6572cad9 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ 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)