在Optuna插件中存储用户属性

Question

当通过九头蛇扫描插件使用optuna trial时，我如何在它中存储额外的信息？

我的用例如下:我想优化一堆超参数。我将所有实验(即试验)的可再现性信息存储在一个单独的数据库中。我知道我可以通过optuna.load_study().best_params甚至best_trial获得最好的价值。然而，这只允许我复制这个实验--这可能需要相当长的时间。为了克服这个问题，我需要以某种方式将它链接到我自己的数据库。我想将自己数据库的ID存储在trial对象中的某个地方。

如果不使用九头蛇，我想我会设置用户属性。然而，与九头蛇把这一切抽象化，似乎没有选择这样做。

我知道，我只需查询我自己的数据库，找出optuna找到的最佳配角的确切组合，但这似乎是解决一个简单问题的一个困难的解决方案。

一些最低限度的代码：

from dataclasses import dataclass

import hydra
from hydra.core.config_store import ConfigStore
from omegaconf import MISSING


@dataclass
class TrainConfig:
    x: float | int = MISSING
    y: int = MISSING
    z: int | None = None


ConfigStore.instance().store(name="config", node=TrainConfig)


@hydra.main(version_base=None, config_path="conf", config_name="sweep")
def sphere(cfg: TrainConfig) -> float:
    x: float = cfg.x
    y: float = cfg.y
    return x**2 + y**2


if __name__ == "__main__":
    sphere()

defaults:
  - override hydra/sweeper: optuna
  - override hydra/sweeper/sampler: tpe

hydra:
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    study_name: sphere
    storage: sqlite:///trials.db
    n_trials: 20
    n_jobs: 1
    params:
      x: range(-5.5, 5.5, step=0.5)
      y: choice(-5 ,0 ,5)
      z: choice(0, 3, 5)

x: 1
y: 1
z: 1

Answer 1

通过custom_search_space解决问题的方法。

hydra:
  sweeper:
    sampler:
      seed: 123
    direction: minimize
    study_name: sphere
    storage: sqlite:///trials.db
    n_trials: 20
    n_jobs: 1
    params:
      x: range(-5.5, 5.5, step=0.5)
      y: choice(-5 ,0 ,5)
      z: choice([0, 1], [2, 3], [2, 5])
    custom_search_space: package.run.configure

def configure(_, trial: Trial) -> None:
    trial.set_user_attr("experiment_db_id", 123456)