From 59eb9db6a64352fc19836361321b61523f323000 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Thu, 10 Aug 2023 16:14:56 +0900 Subject: [PATCH 01/31] Preferential GP --- .../preferential/samplers/__init__.py | 5 + optuna_dashboard/preferential/samplers/_gp.py | 390 ++++++++++++++++++ 2 files changed, 395 insertions(+) create mode 100644 optuna_dashboard/preferential/samplers/__init__.py create mode 100644 optuna_dashboard/preferential/samplers/_gp.py diff --git a/optuna_dashboard/preferential/samplers/__init__.py b/optuna_dashboard/preferential/samplers/__init__.py new file mode 100644 index 00000000..3ec05149 --- /dev/null +++ b/optuna_dashboard/preferential/samplers/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from ._gp import PreferentialGPSampler + +__all__ = ["PreferentialGPSampler"] \ No newline at end of file diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py new file mode 100644 index 00000000..c77d5dc6 --- /dev/null +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -0,0 +1,390 @@ +from __future__ import annotations + +import math +from math import erfc +import random +from typing import Any + +from botorch.acquisition.analytic import LogExpectedImprovement +from botorch.models.gpytorch import GPyTorchModel +from botorch.optim import optimize_acqf +import gpytorch.constraints +import gpytorch.kernels +import gpytorch.likelihoods.gaussian_likelihood +from gpytorch.likelihoods.gaussian_likelihood import GaussianLikelihood +from gpytorch.likelihoods.gaussian_likelihood import Interval +from gpytorch.likelihoods.gaussian_likelihood import Prior +from gpytorch.models.exact_gp import ExactGP +import gpytorch.module +from linear_operator.operators import DiagLinearOperator +from linear_operator.operators import LinearOperator +import numpy as np +import optuna +from optuna import distributions +from optuna import Study +from optuna._transform import _SearchSpaceTransform +from optuna.distributions import BaseDistribution +from optuna.search_space import IntersectionSearchSpace +from optuna.trial import FrozenTrial +import pyro +import pyro.infer.autoguide +import pyro.infer.mcmc +from scipy.special import erfcinv +import torch +from torch import Tensor + +from optuna_preferential._preferences import get_preferences + + +class WeightedGaussianLikelihood(GaussianLikelihood): + def __init__( + self, + weights: torch.Tensor | None = None, + noise_prior: Prior | None = None, + noise_constraint: Interval | None = None, + batch_shape: torch.Size = torch.Size(), + **kwargs: Any, + ) -> None: + super().__init__( + noise_prior=noise_prior, + noise_constraint=noise_constraint, + batch_shape=batch_shape, + **kwargs, + ) + self.weights = weights + + def _shaped_noise_covar( + self, base_shape: torch.Size, *params: Any, **kwargs: Any + ) -> Tensor | LinearOperator: + assert self.weights is not None + assert base_shape[-1] == self.weights.shape[-1] + return DiagLinearOperator(1.0 / self.weights) * super()._shaped_noise_covar( + base_shape, *params, **kwargs + ) + + +def _sample_y( + preferences: np.ndarray, + cov_X_X: np.ndarray, + obs_noise_var: float, + cycles: int, + initial_sample: np.ndarray, +) -> np.ndarray: + # TODO: Refactor and write tests for this function. + + N = cov_X_X.shape[0] + M = len(preferences) + cov_X_X = cov_X_X + np.eye(N) * 1e-6 # Add jitter + cov_X_X_chol = np.linalg.cholesky(cov_X_X) + cov_X_X_inv = np.linalg.inv(cov_X_X) + + # (sI + A K A^T)^-1 = s^-1 I - s^-2 A(K^-1 + s^-1 A^T A)^-1 A^T + + schur = cov_X_X_inv.copy() + np.add.at(schur, (preferences[:, 0], preferences[:, 0]), 1.0 / (2 * obs_noise_var)) + np.add.at(schur, (preferences[:, 1], preferences[:, 1]), 1.0 / (2 * obs_noise_var)) + np.add.at(schur, (preferences[:, 0], preferences[:, 1]), -1.0 / (2 * obs_noise_var)) + np.add.at(schur, (preferences[:, 1], preferences[:, 0]), -1.0 / (2 * obs_noise_var)) + idx_M = np.arange(M) + + schur_inv = np.linalg.inv(schur) + + cov_diff_inv = schur_inv[:, preferences[:, 0]] - schur_inv[:, preferences[:, 1]] + cov_diff_inv = cov_diff_inv[preferences[:, 0], :] - cov_diff_inv[preferences[:, 1], :] + cov_diff_inv *= -1 / (2 * obs_noise_var) ** 2 + cov_diff_inv[idx_M, idx_M] += 1.0 / (2 * obs_noise_var) + + diffs = _orthants_MVN_Gibbs_sampling( + cov_diff_inv, + cycles=cycles, + initial_sample=initial_sample[:, 0] - initial_sample[:, 1], + )[-1] + + random_ys = (cov_X_X_chol @ np.random.randn(N))[preferences] + np.sqrt( + obs_noise_var + ) * np.random.randn(M, 2) + errors = diffs - (random_ys[:, 0] - random_ys[:, 1]) + cov_diff_inv_errors = cov_diff_inv @ errors + + AT_cov_diff_inv_errors = np.zeros((N,)) + np.add.at(AT_cov_diff_inv_errors, preferences[:, 0], cov_diff_inv_errors) + np.add.at(AT_cov_diff_inv_errors, preferences[:, 1], -cov_diff_inv_errors) + + return ( + random_ys + + (cov_X_X @ AT_cov_diff_inv_errors)[preferences] + + obs_noise_var * np.array([[1, -1]]) * cov_diff_inv_errors[:, None] + ) + + +_SQRT2 = math.sqrt(2) + + +def _orthants_MVN_Gibbs_sampling( + cov_inv: np.ndarray, + cycles: int, + initial_sample: np.ndarray | None = None, +) -> np.ndarray: + dim = cov_inv.shape[0] + assert cov_inv.shape == (dim, dim) + + if initial_sample is None: + sample_chain = np.zeros(dim) + else: + sample_chain = initial_sample + + conditional_std = 1 / np.sqrt(np.diag(cov_inv)) + + scaled_cov_inv = cov_inv / np.c_[np.diag(cov_inv)] + + out = np.empty((cycles + 1, dim)) + out[0, :] = sample_chain + + for i in range(cycles): + for j in range(dim): + conditional_mean = sample_chain[j] - scaled_cov_inv[j] @ sample_chain + sample_chain[j] = ( + _one_side_trunc_norm_sampling(lower=-conditional_mean / conditional_std[j]) + * conditional_std[j] + + conditional_mean + ) + out[i + 1, :] = sample_chain + + return out + + +def _one_side_trunc_norm_sampling(lower: float) -> float: + return erfcinv(random.random() * erfc(lower / _SQRT2)) * _SQRT2 + + +class PreferentialGP(GPyTorchModel, ExactGP): + _num_outputs = 1 + + def __init__( + self, + kernel: gpytorch.kernels.Kernel, + noise_prior: Prior | None = None, + noise_constraint: Interval | None = None, + ): + GPyTorchModel.__init__(self) + likelihood = WeightedGaussianLikelihood( + noise_prior=noise_prior, noise_constraint=noise_constraint + ) + ExactGP.__init__(self, train_inputs=None, train_targets=None, likelihood=likelihood) + self.covar_module = kernel + + self._last_params: dict[str, torch.Tensor] | None = None + self._last_mcmc_step_size: float | None = None + + def _pyro_model(self, train_x, train_y): + # 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 = 10) -> None: + if len(preferences) == 0: + # Skip actual MCMC computation + self.set_train_data( + inputs=torch.empty((0, X.shape[-1])), + targets=torch.empty((0,)), + strict=False, + ) + self.likelihood.weights = torch.empty((0,)) + else: + dtype = torch.float64 + + cnt = torch.bincount(preferences.reshape(-1)) + mask = cnt > 0 + train_x = X[mask] + weights = cnt[mask] + + assert isinstance(self.likelihood, WeightedGaussianLikelihood) + self.likelihood.weights = weights + + preferences_np = preferences.detach().numpy() + + all_ys_np = np.zeros((len(preferences), 2)) + train_y = torch.zeros( + ( + len( + train_x, + ) + ), + dtype=dtype, + ) + + 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, + ) + warmup_steps = max(0, cycles - 2) + nuts.setup(warmup_steps=warmup_steps, train_x=train_x, train_y=train_y) + + raw_params = self._last_params or nuts.initial_params + for i in range(cycles): + params = { + name: nuts.transforms[name].inv(value) for name, value in raw_params.items() + } + _set_params(self, params) + self.set_train_data(train_x, train_y, strict=False) + all_ys_np = _sample_y( + preferences=preferences_np, + cov_X_X=self.covar_module(train_x).detach().numpy(), + obs_noise_var=float(self.likelihood.noise_covar.noise), + cycles=10, + initial_sample=all_ys_np, + ) + ys_sum_np = np.zeros((len(X),)) + np.add.at(ys_sum_np, preferences_np.reshape(-1), all_ys_np.reshape(-1)) + ys_sum = torch.from_numpy(ys_sum_np) + train_y[:] = ys_sum[mask] / cnt[mask] + nuts.clear_cache() + raw_params = nuts.sample(raw_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) + _set_params(self, params) + + self._last_params = raw_params + self._last_mcmc_step_size = nuts.step_size + nuts.cleanup() + + def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultivariateNormal: + mean_module = gpytorch.means.ZeroMean() + return gpytorch.distributions.MultivariateNormal( + mean_module(x), + self.covar_module(x), + ) + + +def _set_params( + module: gpytorch.Module, params_dict: dict[str, torch.Tensor], memo=None, prefix="" +): + if memo is None: + memo = set() + if hasattr(module, "_priors"): + for name, (prior, closure, setting_closure) in module._priors.items(): + if prior is not None and prior not in memo: + memo.add(prior) + setting_closure(module, params_dict[prefix + ("." if prefix else "") + name]) + + for mname, module_ in module.named_children(): + submodule_prefix = prefix + ("." if prefix else "") + mname + _set_params(module_, params_dict, memo=memo, prefix=submodule_prefix) + + +class PreferentialGPSampler(optuna.samplers.BaseSampler): + def __init__( + self, + *, + kernel: gpytorch.kernels.Kernel | None = None, + noise_prior: Prior | None = None, + independent_sampler: optuna.samplers.BaseSampler | None = None, + seed: int | None = None, + device: torch.device | None = None, + ) -> None: + self._rng = np.random.RandomState(seed=seed) + self._search_space = IntersectionSearchSpace() + + self.kernel = kernel + self.noise_prior = noise_prior + self.independent_sampler = independent_sampler or optuna.samplers.RandomSampler( + seed=self._rng.randint(2**32), + ) + self.device = device or torch.device("cpu") + + self._gp = None + + def reseed_rng(self) -> None: + self._rng.seed() + self.independent_sampler.reseed_rng() + + def infer_relative_search_space( + self, study: Study, trial: FrozenTrial + ) -> dict[str, BaseDistribution]: + return self._search_space.calculate(study) + + def sample_relative( + self, + study: Study, + trial: FrozenTrial, + search_space: dict[str, BaseDistribution], + ) -> dict[str, Any]: + if len(search_space) == 0: + return {} + + preferences = get_preferences(study, deepcopy=False) + if len(preferences) == 0: + return {} + + trans = _SearchSpaceTransform( + search_space, transform_log=True, transform_step=True, transform_0_1=True + ) + dims = len(trans.bounds) + self._gp = self._gp or PreferentialGP( + kernel=self.kernel + or gpytorch.kernels.MaternKernel( + nu=2.5, + ard_num_dims=dims, + lengthscale_prior=gpytorch.priors.GammaPrior(3.0, 6.0), + lengthscale_constraint=gpytorch.constraints.Positive(), + ), + noise_prior=self.noise_prior or gpytorch.priors.GammaPrior(1.1, 2.0), + noise_constraint=gpytorch.constraints.Positive(), + ) + + ids = {} + params = [] + pref_ids = [] + + 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])) + dtype = torch.float64 + + params_torch = torch.tensor(np.array(params), dtype=dtype, device=self.device) + pref_ids_torch = torch.tensor( + np.array(pref_ids), + dtype=torch.int32, + device=self.device, + ) + + self._gp.fit_mcmc(params_torch, pref_ids_torch) + self._gp.eval() + scores = self._gp(params_torch).mean + + best_f = torch.max(scores) + + acqf = LogExpectedImprovement( + model=self._gp, + best_f=best_f, + ) + + # TODO: Make it possible to apply it on categorical variables + candidates, _ = optimize_acqf( + acq_function=acqf, + bounds=torch.from_numpy(trans.bounds.T), + q=1, + num_restarts=10, + raw_samples=512, + options={"batch_limit": 5, "maxiter": 200}, + sequential=True, + ) + next_x = trans.untransform(candidates[0].detach().numpy()) + return next_x + + def sample_independent( + self, + study: Study, + trial: FrozenTrial, + param_name: str, + param_distribution: distributions.BaseDistribution, + ) -> Any: + return self.independent_sampler.sample_independent( + study, trial, param_name, param_distribution + ) From 38407a198cc882492f80c951c299c7e4f5fa79a2 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Thu, 10 Aug 2023 16:22:00 +0900 Subject: [PATCH 02/31] format --- optuna_dashboard/preferential/samplers/__init__.py | 3 ++- optuna_dashboard/preferential/samplers/_gp.py | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/optuna_dashboard/preferential/samplers/__init__.py b/optuna_dashboard/preferential/samplers/__init__.py index 3ec05149..f68708f4 100644 --- a/optuna_dashboard/preferential/samplers/__init__.py +++ b/optuna_dashboard/preferential/samplers/__init__.py @@ -2,4 +2,5 @@ from __future__ import annotations from ._gp import PreferentialGPSampler -__all__ = ["PreferentialGPSampler"] \ No newline at end of file + +__all__ = ["PreferentialGPSampler"] diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index c77d5dc6..706a2327 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -26,6 +26,7 @@ from optuna._transform import _SearchSpaceTransform from optuna.distributions import BaseDistribution from optuna.search_space import IntersectionSearchSpace from optuna.trial import FrozenTrial +from optuna_preferential._preferences import get_preferences import pyro import pyro.infer.autoguide import pyro.infer.mcmc @@ -33,8 +34,6 @@ from scipy.special import erfcinv import torch from torch import Tensor -from optuna_preferential._preferences import get_preferences - class WeightedGaussianLikelihood(GaussianLikelihood): def __init__( From 1ca3c07e9d224fbf77be0990815427f1a26ec631 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Thu, 10 Aug 2023 16:29:27 +0900 Subject: [PATCH 03/31] mypy --- optuna_dashboard/preferential/samplers/_gp.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index 706a2327..e5228dc7 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -26,7 +26,6 @@ from optuna._transform import _SearchSpaceTransform from optuna.distributions import BaseDistribution from optuna.search_space import IntersectionSearchSpace from optuna.trial import FrozenTrial -from optuna_preferential._preferences import get_preferences import pyro import pyro.infer.autoguide import pyro.infer.mcmc @@ -34,8 +33,10 @@ from scipy.special import erfcinv import torch from torch import Tensor +from .._preferences import get_preferences -class WeightedGaussianLikelihood(GaussianLikelihood): + +class _WeightedGaussianLikelihood(GaussianLikelihood): def __init__( self, weights: torch.Tensor | None = None, @@ -156,7 +157,7 @@ def _one_side_trunc_norm_sampling(lower: float) -> float: return erfcinv(random.random() * erfc(lower / _SQRT2)) * _SQRT2 -class PreferentialGP(GPyTorchModel, ExactGP): +class _PreferentialGP(GPyTorchModel, ExactGP): _num_outputs = 1 def __init__( @@ -164,9 +165,9 @@ class PreferentialGP(GPyTorchModel, ExactGP): kernel: gpytorch.kernels.Kernel, noise_prior: Prior | None = None, noise_constraint: Interval | None = None, - ): + ) -> None: GPyTorchModel.__init__(self) - likelihood = WeightedGaussianLikelihood( + likelihood = _WeightedGaussianLikelihood( noise_prior=noise_prior, noise_constraint=noise_constraint ) ExactGP.__init__(self, train_inputs=None, train_targets=None, likelihood=likelihood) @@ -175,7 +176,7 @@ class PreferentialGP(GPyTorchModel, ExactGP): self._last_params: dict[str, torch.Tensor] | None = None self._last_mcmc_step_size: float | None = None - def _pyro_model(self, train_x, train_y): + 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)) @@ -259,8 +260,8 @@ class PreferentialGP(GPyTorchModel, ExactGP): def _set_params( - module: gpytorch.Module, params_dict: dict[str, torch.Tensor], memo=None, prefix="" -): + module: gpytorch.Module, params_dict: dict[str, torch.Tensor], memo: set | None=None, prefix: str="" +) -> None: if memo is None: memo = set() if hasattr(module, "_priors"): @@ -294,7 +295,7 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): ) self.device = device or torch.device("cpu") - self._gp = None + self._gp: PreferentialGPSampler | None = None def reseed_rng(self) -> None: self._rng.seed() @@ -322,7 +323,7 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): search_space, transform_log=True, transform_step=True, transform_0_1=True ) dims = len(trans.bounds) - self._gp = self._gp or PreferentialGP( + self._gp = self._gp or _PreferentialGP( kernel=self.kernel or gpytorch.kernels.MaternKernel( nu=2.5, From 0c35254100bbbea108927fb11655707a1c8bfe26 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Thu, 10 Aug 2023 16:49:59 +0900 Subject: [PATCH 04/31] fix typo --- optuna_dashboard/preferential/samplers/_gp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index e5228dc7..5eee4a78 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -199,7 +199,7 @@ class _PreferentialGP(GPyTorchModel, ExactGP): train_x = X[mask] weights = cnt[mask] - assert isinstance(self.likelihood, WeightedGaussianLikelihood) + assert isinstance(self.likelihood, _WeightedGaussianLikelihood) self.likelihood.weights = weights preferences_np = preferences.detach().numpy() From bb9a42b843fbcd1a0a4f57ec492af660f1e03c77 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 11:08:37 +0900 Subject: [PATCH 05/31] Add examples --- .../preferential-optimization/evaluator.py | 93 +++++++++++++++++++ .../preferential-optimization/generator.py | 63 +++++++++++++ 2 files changed, 156 insertions(+) create mode 100644 examples/preferential-optimization/evaluator.py create mode 100644 examples/preferential-optimization/generator.py diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py new file mode 100644 index 00000000..c03f598b --- /dev/null +++ b/examples/preferential-optimization/evaluator.py @@ -0,0 +1,93 @@ +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 +import streamlit as st + +from optuna_dashboard.preferential import load_study + + +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. 比較対象のTrialを取得 + comparison_trials = study.best_trials + + st.text("Which is the worst?") + + # 2. 各TrialのArtifact画像を並べて表示 + 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.get("rgb_artifact_id") + image_caption = trial.user_attrs.get("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: + time.sleep(0.1) + st.experimental_rerun() + + +if __name__ == "__main__": + main() diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py new file mode 100644 index 00000000..d68f6c69 --- /dev/null +++ b/examples/preferential-optimization/generator.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import os +import tempfile +import time +from time import sleep +from typing import NoReturn + +import optuna +from optuna_dashboard.artifact import upload_artifact +from optuna_dashboard.artifact.file_system import FileSystemBackend +from PIL import Image + +from optuna_dashboard.preferential import create_study +from optuna_dashboard.preferential.samplers._gp import PreferentialGPSampler + + +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 main() -> NoReturn: + study = create_study( + study_name="Preferential Optimization", + storage=STORAGE_URL, + sampler=PreferentialGPSampler(), + load_if_exists=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + while True: + if len(study.best_trials) >= n_comparison: + time.sleep(0.1) # Avoid busy-loop + continue + + trial = study.ask() + # 1. Ask new parameters + r = trial.suggest_int("r", 0, 255) + g = trial.suggest_int("g", 0, 255) + b = trial.suggest_int("b", 0, 255) + + # 2. Generate image + image_path = os.path.join(tmpdir, f"sample-{trial.number}.png") + image = Image.new("RGB", (320, 240), color=(r, g, b)) + # sleep(2.0) + image.save(image_path) + + # 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})") + print("RGB:", (r, g, b)) + + # 4. Mark comparison ready + study.mark_comparison_ready(trial) + + +if __name__ == "__main__": + main() From bcf29d5983ca1cdfcc0b279c9f9c77a54a788c45 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 11:57:07 +0900 Subject: [PATCH 06/31] Add note --- examples/preferential-optimization/evaluator.py | 1 + examples/preferential-optimization/generator.py | 13 +++++++++++++ optuna_dashboard/preferential/samplers/_gp.py | 2 +- 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index c03f598b..dc660691 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -12,6 +12,7 @@ from optuna_dashboard.artifact.file_system import FileSystemBackend import streamlit as st from optuna_dashboard.preferential import load_study +from optuna_dashboard. STORAGE_URL = "sqlite:///st-example.db" diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index d68f6c69..659505f5 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -14,6 +14,9 @@ from PIL import Image from optuna_dashboard.preferential import create_study from optuna_dashboard.preferential.samplers._gp import PreferentialGPSampler +from optuna_dashboard import save_note +import textwrap + STORAGE_URL = "sqlite:///st-example.db" artifact_path = os.path.join(os.path.dirname(__file__), "artifact") @@ -55,6 +58,16 @@ def main() -> NoReturn: trial.set_user_attr("image_caption", f"(R, G, B) = ({r}, {g}, {b})") print("RGB:", (r, g, b)) + + # 4. Save Note + note = textwrap.dedent( + f"""\ + ![generated-image]({artifact_path}) + (R, G, B) = ({r}, {g}, {b}) + """ + ) + save_note(trial, note) + # 4. Mark comparison ready study.mark_comparison_ready(trial) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index 5eee4a78..9b62f98d 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -33,7 +33,7 @@ from scipy.special import erfcinv import torch from torch import Tensor -from .._preferences import get_preferences +from .._system_attrs import get_preferences class _WeightedGaussianLikelihood(GaussianLikelihood): From 7c3cc3e1adf1455ddfbabaae79f05ad44cb8ab84 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 11:57:32 +0900 Subject: [PATCH 07/31] Change comment --- examples/preferential-optimization/generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index 659505f5..5b9aab43 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -68,7 +68,7 @@ def main() -> NoReturn: ) save_note(trial, note) - # 4. Mark comparison ready + # 5. Mark comparison ready study.mark_comparison_ready(trial) From f8d7bf11d9ba11e299a30b07b5a45e5be576c7bc Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 12:00:10 +0900 Subject: [PATCH 08/31] Revert evaluator --- examples/preferential-optimization/evaluator.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index dc660691..c03f598b 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -12,7 +12,6 @@ from optuna_dashboard.artifact.file_system import FileSystemBackend import streamlit as st from optuna_dashboard.preferential import load_study -from optuna_dashboard. STORAGE_URL = "sqlite:///st-example.db" From c23fda70492f2b685c0b923600ba1eb6b2808045 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 16:43:10 +0900 Subject: [PATCH 09/31] format --- examples/preferential-optimization/evaluator.py | 4 ++-- examples/preferential-optimization/generator.py | 11 +++-------- optuna_dashboard/preferential/samplers/_gp.py | 5 ++++- 3 files changed, 9 insertions(+), 11 deletions(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index c03f598b..1909ddd8 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -9,10 +9,10 @@ from typing import NoReturn import uuid from optuna_dashboard.artifact.file_system import FileSystemBackend -import streamlit as st - 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") diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index 5b9aab43..63e81183 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -2,20 +2,16 @@ from __future__ import annotations import os import tempfile +import textwrap import time -from time import sleep from typing import NoReturn -import optuna +from optuna_dashboard import save_note from optuna_dashboard.artifact import upload_artifact from optuna_dashboard.artifact.file_system import FileSystemBackend -from PIL import Image - from optuna_dashboard.preferential import create_study from optuna_dashboard.preferential.samplers._gp import PreferentialGPSampler - -from optuna_dashboard import save_note -import textwrap +from PIL import Image STORAGE_URL = "sqlite:///st-example.db" @@ -58,7 +54,6 @@ def main() -> NoReturn: trial.set_user_attr("image_caption", f"(R, G, B) = ({r}, {g}, {b})") print("RGB:", (r, g, b)) - # 4. Save Note note = textwrap.dedent( f"""\ diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index 9b62f98d..86192a78 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -260,7 +260,10 @@ class _PreferentialGP(GPyTorchModel, ExactGP): def _set_params( - module: gpytorch.Module, params_dict: dict[str, torch.Tensor], memo: set | None=None, prefix: str="" + module: gpytorch.Module, + params_dict: dict[str, torch.Tensor], + memo: set | None = None, + prefix: str = "", ) -> None: if memo is None: memo = set() From b7f90fe6b75f6553810c8f4d58a2c281a1027d5a Mon Sep 17 00:00:00 2001 From: Contramundum Date: Mon, 14 Aug 2023 17:22:37 +0900 Subject: [PATCH 10/31] mypy --- optuna_dashboard/preferential/samplers/_gp.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index 86192a78..d6d4caff 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -298,7 +298,7 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): ) self.device = device or torch.device("cpu") - self._gp: PreferentialGPSampler | None = None + self._gp: _PreferentialGP | None = None def reseed_rng(self) -> None: self._rng.seed() @@ -338,9 +338,9 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): noise_constraint=gpytorch.constraints.Positive(), ) - ids = {} - params = [] - pref_ids = [] + ids: dict[int, int] = {} + params: list[torch.Tensor] = [] + pref_ids: list[tuple[int, int]] = [] for better, worse in preferences: for t in (better, worse): From d01dfd013ef06ca95a02941bd7da3d8a9c1935bc Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 16:34:30 +0900 Subject: [PATCH 11/31] Update examples/preferential-optimization/evaluator.py --- examples/preferential-optimization/evaluator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index 1909ddd8..d14f442f 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -39,7 +39,7 @@ def main() -> NoReturn: storage=STORAGE_URL, ) - # 1. 比較対象のTrialを取得 + # 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?") From c66de309c12a0ab7807e74d80824aa495bea03ea Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 16:38:49 +0900 Subject: [PATCH 12/31] Update examples/preferential-optimization/evaluator.py --- examples/preferential-optimization/evaluator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index d14f442f..71a1df9d 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -44,7 +44,7 @@ def main() -> NoReturn: st.text("Which is the worst?") - # 2. 各TrialのArtifact画像を並べて表示 + # 2. Show the artifact images of all those trials. cols = st.columns(n_comparison) finished_dict = {t.number: t for t in comparison_trials} From 0a6cb1d9ffc5a2efa4711e54ddbda24050b82c9e Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 16:49:23 +0900 Subject: [PATCH 13/31] Update examples/preferential-optimization/evaluator.py --- examples/preferential-optimization/evaluator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index 71a1df9d..355af4d7 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -70,8 +70,8 @@ def main() -> NoReturn: trial = finished_dict[trial_number] col = cols[col_i] - rgb_artifact_id = trial.user_attrs.get("rgb_artifact_id") - image_caption = trial.user_attrs.get("image_caption") + 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") From a39eaecc9af2be73b8fb3e372f69bdd8b90cbf0f Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 16:50:30 +0900 Subject: [PATCH 14/31] Update examples/preferential-optimization/evaluator.py --- examples/preferential-optimization/evaluator.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/preferential-optimization/evaluator.py b/examples/preferential-optimization/evaluator.py index 355af4d7..ee1950c5 100644 --- a/examples/preferential-optimization/evaluator.py +++ b/examples/preferential-optimization/evaluator.py @@ -85,6 +85,7 @@ def main() -> NoReturn: continue if len(comparison_trials) < n_comparison: + # Wait for unfinished trials (images under generation) to be generated. time.sleep(0.1) st.experimental_rerun() From 6a9c05f93f1da0953ccb4615a933dc6c483f0c19 Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 17:05:15 +0900 Subject: [PATCH 15/31] Update examples/preferential-optimization/generator.py --- examples/preferential-optimization/generator.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index 63e81183..87773bd7 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -32,6 +32,8 @@ def main() -> NoReturn: with tempfile.TemporaryDirectory() as tmpdir: while True: + # If n_comparison "best" trials (that are not reported bad) exists, + # the generator waits for human evaluation. if len(study.best_trials) >= n_comparison: time.sleep(0.1) # Avoid busy-loop continue From 941e5011d5ddabc8cc7c0df8c397c130807ba044 Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Wed, 16 Aug 2023 17:05:39 +0900 Subject: [PATCH 16/31] Update examples/preferential-optimization/generator.py --- examples/preferential-optimization/generator.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index 87773bd7..07977d0b 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -47,7 +47,6 @@ def main() -> NoReturn: # 2. Generate image image_path = os.path.join(tmpdir, f"sample-{trial.number}.png") image = Image.new("RGB", (320, 240), color=(r, g, b)) - # sleep(2.0) image.save(image_path) # 3. Upload Artifact From b30435d6776ebbb1fa91b04f5092a310487a6e00 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Thu, 17 Aug 2023 15:55:07 +0900 Subject: [PATCH 17/31] Make it possible to fix random seed --- optuna_dashboard/preferential/samplers/_gp.py | 143 +++++++++--------- 1 file changed, 75 insertions(+), 68 deletions(-) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/_gp.py index d6d4caff..39d85b5b 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -2,7 +2,6 @@ from __future__ import annotations import math from math import erfc -import random from typing import Any from botorch.acquisition.analytic import LogExpectedImprovement @@ -69,6 +68,7 @@ def _sample_y( obs_noise_var: float, cycles: int, initial_sample: np.ndarray, + rng: np.random.RandomState ) -> np.ndarray: # TODO: Refactor and write tests for this function. @@ -98,11 +98,12 @@ def _sample_y( cov_diff_inv, cycles=cycles, initial_sample=initial_sample[:, 0] - initial_sample[:, 1], + rng=rng, )[-1] - random_ys = (cov_X_X_chol @ np.random.randn(N))[preferences] + np.sqrt( + random_ys = (cov_X_X_chol @ rng.randn(N))[preferences] + np.sqrt( obs_noise_var - ) * np.random.randn(M, 2) + ) * rng.randn(M, 2) errors = diffs - (random_ys[:, 0] - random_ys[:, 1]) cov_diff_inv_errors = cov_diff_inv @ errors @@ -123,7 +124,8 @@ _SQRT2 = math.sqrt(2) def _orthants_MVN_Gibbs_sampling( cov_inv: np.ndarray, cycles: int, - initial_sample: np.ndarray | None = None, + initial_sample: np.ndarray, + rng: np.random.RandomState, ) -> np.ndarray: dim = cov_inv.shape[0] assert cov_inv.shape == (dim, dim) @@ -144,7 +146,7 @@ def _orthants_MVN_Gibbs_sampling( for j in range(dim): conditional_mean = sample_chain[j] - scaled_cov_inv[j] @ sample_chain sample_chain[j] = ( - _one_side_trunc_norm_sampling(lower=-conditional_mean / conditional_std[j]) + _one_side_trunc_norm_sampling(lower=-conditional_mean / conditional_std[j], rng=rng) * conditional_std[j] + conditional_mean ) @@ -153,8 +155,8 @@ def _orthants_MVN_Gibbs_sampling( return out -def _one_side_trunc_norm_sampling(lower: float) -> float: - return erfcinv(random.random() * erfc(lower / _SQRT2)) * _SQRT2 +def _one_side_trunc_norm_sampling(lower: float, rng: np.random.RandomState) -> float: + return erfcinv(rng.rand() * erfc(lower / _SQRT2)) * _SQRT2 class _PreferentialGP(GPyTorchModel, ExactGP): @@ -182,7 +184,7 @@ class _PreferentialGP(GPyTorchModel, ExactGP): 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 = 10) -> None: + def fit_mcmc(self, X: torch.Tensor, preferences: torch.Tensor, cycles: int, rng: np.random.RandomState) -> None: if len(preferences) == 0: # Skip actual MCMC computation self.set_train_data( @@ -235,6 +237,7 @@ class _PreferentialGP(GPyTorchModel, ExactGP): obs_noise_var=float(self.likelihood.noise_covar.noise), cycles=10, initial_sample=all_ys_np, + rng=rng, ) ys_sum_np = np.zeros((len(X),)) np.add.at(ys_sum_np, preferences_np.reshape(-1), all_ys_np.reshape(-1)) @@ -288,7 +291,7 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): seed: int | None = None, device: torch.device | None = None, ) -> None: - self._rng = np.random.RandomState(seed=seed) + self._rng = np.random.RandomState(seed) self._search_space = IntersectionSearchSpace() self.kernel = kernel @@ -300,9 +303,10 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): self._gp: _PreferentialGP | None = None + def reseed_rng(self) -> None: - self._rng.seed() self.independent_sampler.reseed_rng() + self._rng = np.random.RandomState() def infer_relative_search_space( self, study: Study, trial: FrozenTrial @@ -315,71 +319,74 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): trial: FrozenTrial, search_space: dict[str, BaseDistribution], ) -> dict[str, Any]: - if len(search_space) == 0: - return {} + with torch.random.fork_rng(): + torch.manual_seed(self._rng.randint(2 ** 32)) + pyro.set_rng_seed(self._rng.randint(2 ** 32)) + + if len(search_space) == 0: + return {} + + preferences = get_preferences(study, deepcopy=False) + if len(preferences) == 0: + return {} - preferences = get_preferences(study, deepcopy=False) - if len(preferences) == 0: - return {} + trans = _SearchSpaceTransform( + search_space, transform_log=True, transform_step=True, transform_0_1=True + ) + dims = len(trans.bounds) + self._gp = self._gp or _PreferentialGP( + kernel=self.kernel + or gpytorch.kernels.MaternKernel( + nu=2.5, + ard_num_dims=dims, + lengthscale_prior=gpytorch.priors.GammaPrior(3.0, 6.0), + lengthscale_constraint=gpytorch.constraints.Positive(), + ), + noise_prior=self.noise_prior or gpytorch.priors.GammaPrior(1.1, 2.0), + noise_constraint=gpytorch.constraints.Positive(), + ) - trans = _SearchSpaceTransform( - search_space, transform_log=True, transform_step=True, transform_0_1=True - ) - dims = len(trans.bounds) - self._gp = self._gp or _PreferentialGP( - kernel=self.kernel - or gpytorch.kernels.MaternKernel( - nu=2.5, - ard_num_dims=dims, - lengthscale_prior=gpytorch.priors.GammaPrior(3.0, 6.0), - lengthscale_constraint=gpytorch.constraints.Positive(), - ), - noise_prior=self.noise_prior or gpytorch.priors.GammaPrior(1.1, 2.0), - noise_constraint=gpytorch.constraints.Positive(), - ) + ids: dict[int, int] = {} + params: list[torch.Tensor] = [] + pref_ids: list[tuple[int, int]] = [] - ids: dict[int, int] = {} - params: list[torch.Tensor] = [] - pref_ids: list[tuple[int, int]] = [] + 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])) + dtype = torch.float64 - 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])) - dtype = torch.float64 + params_torch = torch.tensor(np.array(params), dtype=dtype, device=self.device) + pref_ids_torch = torch.tensor( + np.array(pref_ids), + dtype=torch.int32, + device=self.device, + ) + self._gp.fit_mcmc(params_torch, pref_ids_torch, cycles=10, rng=self._rng) + self._gp.eval() + scores = self._gp(params_torch).mean - params_torch = torch.tensor(np.array(params), dtype=dtype, device=self.device) - pref_ids_torch = torch.tensor( - np.array(pref_ids), - dtype=torch.int32, - device=self.device, - ) + best_f = torch.max(scores) - self._gp.fit_mcmc(params_torch, pref_ids_torch) - self._gp.eval() - scores = self._gp(params_torch).mean + acqf = LogExpectedImprovement( + model=self._gp, + best_f=best_f, + ) - best_f = torch.max(scores) - - acqf = LogExpectedImprovement( - model=self._gp, - best_f=best_f, - ) - - # TODO: Make it possible to apply it on categorical variables - candidates, _ = optimize_acqf( - acq_function=acqf, - bounds=torch.from_numpy(trans.bounds.T), - q=1, - num_restarts=10, - raw_samples=512, - options={"batch_limit": 5, "maxiter": 200}, - sequential=True, - ) - next_x = trans.untransform(candidates[0].detach().numpy()) - return next_x + # TODO: Make it possible to apply it on categorical variables + candidates, _ = optimize_acqf( + acq_function=acqf, + bounds=torch.from_numpy(trans.bounds.T), + q=1, + num_restarts=10, + raw_samples=512, + options={"batch_limit": 5, "maxiter": 200}, + sequential=True, + ) + next_x = trans.untransform(candidates[0].detach().numpy()) + return next_x def sample_independent( self, From 744bdee4de35b1408ed8c80ccec39c499c6a5253 Mon Sep 17 00:00:00 2001 From: Contramundum Date: Tue, 22 Aug 2023 19:27:25 +0900 Subject: [PATCH 18/31] 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) From 140ca3957b4670cd18bbec1e3c632cf967c94acf Mon Sep 17 00:00:00 2001 From: Contramundum Date: Tue, 22 Aug 2023 19:28:52 +0900 Subject: [PATCH 19/31] format --- .../preferential-optimization/generator.py | 5 +-- optuna_dashboard/preferential/samplers/_gp.py | 33 ++++++++++--------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index bbfea678..769ad96b 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -7,7 +7,8 @@ import time from typing import NoReturn from optuna_dashboard import save_note -from optuna_dashboard.artifact import upload_artifact, get_artifact_path +from optuna_dashboard.artifact import get_artifact_path +from optuna_dashboard.artifact import upload_artifact from optuna_dashboard.artifact.file_system import FileSystemBackend from optuna_dashboard.preferential import create_study from optuna_dashboard.preferential.samplers._gp import PreferentialGPSampler @@ -58,7 +59,7 @@ def main() -> NoReturn: note = textwrap.dedent( f"""\ ![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 6572cad9..ff4001c1 100644 --- a/optuna_dashboard/preferential/samplers/_gp.py +++ b/optuna_dashboard/preferential/samplers/_gp.py @@ -17,6 +17,7 @@ from gpytorch.models.exact_gp import ExactGP import gpytorch.module from linear_operator.operators import DiagLinearOperator from linear_operator.operators import LinearOperator +from linear_operator.utils.errors import NotPSDError import numpy as np import optuna from optuna import distributions @@ -31,7 +32,6 @@ 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 @@ -69,7 +69,7 @@ def _sample_y( obs_noise_var: float, cycles: int, initial_sample: np.ndarray, - rng: np.random.RandomState + rng: np.random.RandomState, ) -> np.ndarray: # TODO: Refactor and write tests for this function. @@ -102,9 +102,9 @@ def _sample_y( rng=rng, )[-1] - random_ys = (cov_X_X_chol @ rng.randn(N))[preferences] + np.sqrt( - obs_noise_var - ) * rng.randn(M, 2) + random_ys = (cov_X_X_chol @ rng.randn(N))[preferences] + np.sqrt(obs_noise_var) * rng.randn( + M, 2 + ) errors = diffs - (random_ys[:, 0] - random_ys[:, 1]) cov_diff_inv_errors = cov_diff_inv @ errors @@ -147,7 +147,9 @@ def _orthants_MVN_Gibbs_sampling( for j in range(dim): conditional_mean = sample_chain[j] - scaled_cov_inv[j] @ sample_chain sample_chain[j] = ( - _one_side_trunc_norm_sampling(lower=-conditional_mean / conditional_std[j], rng=rng) + _one_side_trunc_norm_sampling( + lower=-conditional_mean / conditional_std[j], rng=rng + ) * conditional_std[j] + conditional_mean ) @@ -182,12 +184,14 @@ 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: + def fit_mcmc( + self, X: torch.Tensor, preferences: torch.Tensor, cycles: int, rng: np.random.RandomState + ) -> None: if len(preferences) == 0: # Skip actual MCMC computation self.set_train_data( @@ -259,8 +263,6 @@ class _PreferentialGP(GPyTorchModel, ExactGP): 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) _set_params(self, params) @@ -318,7 +320,6 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): self._gp: _PreferentialGP | None = None - def reseed_rng(self) -> None: self.independent_sampler.reseed_rng() self._rng = np.random.RandomState() @@ -335,12 +336,12 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): search_space: dict[str, BaseDistribution], ) -> dict[str, Any]: with torch.random.fork_rng(): - torch.manual_seed(self._rng.randint(2 ** 32)) - pyro.set_rng_seed(self._rng.randint(2 ** 32)) - + torch.manual_seed(self._rng.randint(2**32)) + pyro.set_rng_seed(self._rng.randint(2**32)) + if len(search_space) == 0: return {} - + preferences = get_preferences(study._study_id, study._storage) trials = study.get_trials(deepcopy=False) if len(preferences) == 0: From 32740710ea08b97387fae5082ceb809e72d9c94d Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Fri, 25 Aug 2023 20:13:28 +0900 Subject: [PATCH 20/31] Update __init__.py --- optuna_dashboard/preferential/samplers/__init__.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/optuna_dashboard/preferential/samplers/__init__.py b/optuna_dashboard/preferential/samplers/__init__.py index f68708f4..8b137891 100644 --- a/optuna_dashboard/preferential/samplers/__init__.py +++ b/optuna_dashboard/preferential/samplers/__init__.py @@ -1,6 +1 @@ -from __future__ import annotations -from ._gp import PreferentialGPSampler - - -__all__ = ["PreferentialGPSampler"] From ec179af813dd294e152e758e2a96b9bef7f6d0ae Mon Sep 17 00:00:00 2001 From: contramundum53 Date: Fri, 25 Aug 2023 20:13:48 +0900 Subject: [PATCH 21/31] Rename _gp.py to gp.py --- optuna_dashboard/preferential/samplers/{_gp.py => gp.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename optuna_dashboard/preferential/samplers/{_gp.py => gp.py} (100%) diff --git a/optuna_dashboard/preferential/samplers/_gp.py b/optuna_dashboard/preferential/samplers/gp.py similarity index 100% rename from optuna_dashboard/preferential/samplers/_gp.py rename to optuna_dashboard/preferential/samplers/gp.py From 08e44c7abaaeb0fa274ef3b93bc75696d71d7039 Mon Sep 17 00:00:00 2001 From: moririn2528 Date: Mon, 28 Aug 2023 14:26:08 +0900 Subject: [PATCH 22/31] add skip flag --- optuna_dashboard/_app.py | 19 ++++++++++++++++++- optuna_dashboard/preferential/_study.py | 14 +++++++++++--- .../preferential/_system_attrs.py | 17 +++++++++++++++++ optuna_dashboard/ts/action.ts | 12 ++++++++++++ optuna_dashboard/ts/apiClient.ts | 11 +++++++++++ .../ts/components/PreferentialTrials.tsx | 9 +++++++++ 6 files changed, 78 insertions(+), 4 deletions(-) diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 2401f3f0..0ec82fe8 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -40,7 +40,7 @@ from .artifact._backend import register_artifact_route from .artifact._backend_to_store import to_artifact_store from .preferential._study import _SYSTEM_ATTR_PREFERENTIAL_STUDY from .preferential._study import get_best_trials as get_best_preferential_trials -from .preferential._system_attrs import report_preferences +from .preferential._system_attrs import report_preferences, report_skip if typing.TYPE_CHECKING: @@ -331,6 +331,23 @@ def create_app( response.status = 204 return {} + @app.post("/api/studies///skip") + @json_api_view + def skip_trial(study_id: int, trial_id: int) -> dict[str, Any]: + summary = get_study_summary(storage, study_id) + if summary is None: + response.status = 404 # Not found + return {"reason": f"study_id={study_id} is not found"} + system_attrs = getattr(summary, "system_attrs", {}) + is_preferential = system_attrs.get(_SYSTEM_ATTR_PREFERENTIAL_STUDY, False) + if not is_preferential: + response.status = 400 # Bad request + return {"reason": "The study is not preferential."} + + report_skip(study_id, trial_id, storage) + response.status = 204 # No content + return {} + @app.put("/api/studies///note") @json_api_view def save_trial_note(study_id: int, trial_id: int) -> dict[str, Any]: diff --git a/optuna_dashboard/preferential/_study.py b/optuna_dashboard/preferential/_study.py index 271eabfd..d74fbf14 100644 --- a/optuna_dashboard/preferential/_study.py +++ b/optuna_dashboard/preferential/_study.py @@ -12,8 +12,11 @@ from optuna.samplers import BaseSampler from optuna.samplers import RandomSampler from optuna.trial import FrozenTrial from optuna.trial import TrialState -from optuna_dashboard.preferential._system_attrs import get_preferences -from optuna_dashboard.preferential._system_attrs import report_preferences +from optuna_dashboard.preferential._system_attrs import ( + get_preferences, + report_preferences, + _SYSTEM_ATTR_SKIP_TRIAL, +) _logger = logging.get_logger(__name__) @@ -265,7 +268,12 @@ def get_best_trials(study_id: int, storage: optuna.storages.BaseStorage) -> list ] preferences = get_preferences(study_id, storage) worse_numbers = {worse for _, worse in preferences} - return [copy.deepcopy(t) for t in ready_trials if t.number not in worse_numbers] + skiped_numbers = storage.get_study_system_attrs(study_id).get(_SYSTEM_ATTR_SKIP_TRIAL, []) + return [ + copy.deepcopy(t) + for t in ready_trials + if t.number not in worse_numbers and t.number not in skiped_numbers + ] def create_study( diff --git a/optuna_dashboard/preferential/_system_attrs.py b/optuna_dashboard/preferential/_system_attrs.py index fdd9db35..0c2c8777 100644 --- a/optuna_dashboard/preferential/_system_attrs.py +++ b/optuna_dashboard/preferential/_system_attrs.py @@ -9,6 +9,7 @@ from .._storage import get_study_summary _SYSTEM_ATTR_PREFIX_PREFERENCE = "preference:values" +_SYSTEM_ATTR_SKIP_TRIAL = "preference:skip_trial" def report_preferences( @@ -44,3 +45,19 @@ def get_preferences( continue preferences.extend(v) # type: ignore return preferences + + +def report_skip( + study_id: int, + trial_id: int, + storage: BaseStorage, +): + system_attrs = storage.get_study_system_attrs(study_id) + trial_skiped: list[int] = system_attrs.get(_SYSTEM_ATTR_SKIP_TRIAL, []) + trial_number = storage.get_trial(trial_id).number + trial_skiped.append(trial_number) + storage.set_study_system_attr( + study_id=study_id, + key=_SYSTEM_ATTR_SKIP_TRIAL, + value=trial_skiped, + ) diff --git a/optuna_dashboard/ts/action.ts b/optuna_dashboard/ts/action.ts index c0fed1aa..017751fa 100644 --- a/optuna_dashboard/ts/action.ts +++ b/optuna_dashboard/ts/action.ts @@ -15,6 +15,7 @@ import { getMetaInfoAPI, deleteArtifactAPI, reportPreferenceAPI, + skipPreferentialTrialAPI, } from "./apiClient" import { graphVisibilityState, @@ -598,6 +599,16 @@ export const actionCreator = () => { }) } + const skipPreferentialTrial = (studyId: number, trialId: number) => { + skipPreferentialTrialAPI(studyId, trialId).catch((err) => { + const reason = err.response?.data.reason + enqueueSnackbar(`Failed to skip trial. Reason: ${reason}`, { + variant: "error", + }) + console.log(err) + }) + } + return { updateAPIMeta, updateStudyDetail, @@ -618,6 +629,7 @@ export const actionCreator = () => { makeTrialFail, saveTrialUserAttrs, updatePreference, + skipPreferentialTrial, } } diff --git a/optuna_dashboard/ts/apiClient.ts b/optuna_dashboard/ts/apiClient.ts index 5068dc3e..25ca3541 100644 --- a/optuna_dashboard/ts/apiClient.ts +++ b/optuna_dashboard/ts/apiClient.ts @@ -326,3 +326,14 @@ export const reportPreferenceAPI = ( return }) } + +export const skipPreferentialTrialAPI = ( + studyId: number, + trialId: number +): Promise => { + return axiosInstance + .post(`/api/studies/${studyId}/${trialId}/skip`) + .then(() => { + return + }) +} diff --git a/optuna_dashboard/ts/components/PreferentialTrials.tsx b/optuna_dashboard/ts/components/PreferentialTrials.tsx index 9b57586d..c6e7eeb0 100644 --- a/optuna_dashboard/ts/components/PreferentialTrials.tsx +++ b/optuna_dashboard/ts/components/PreferentialTrials.tsx @@ -28,6 +28,15 @@ const PreferentialTrial: FC<{ > Trial {trial.number} (trial_id={trial.trial_id}) +