如何在Hydra Optuna Sweeper中存储Optuna trial用户属性

场景说明 #

• 超参数优化过程中，需要为每个实验trial关联独立的可复现性信息数据库记录，计划将自有数据库的对应记录ID直接存入trial对象，避免后续通过最优参数反向匹配记录的繁琐操作。
• 原生Optuna提供了User Attributes接口支持自定义属性写入，但Hydra对Optuna的底层执行逻辑做了封装，默认没有暴露直接操作当前trial对象的公开入口。

最小复现代码 #

主程序（train.py） #

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()

Sweep配置（conf/sweep.yaml） #

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

Hydra的Optuna Sweeper在调度每个trial执行目标函数时，会将当前trial实例绑定到内置的_TrialCallback回调对象中，并存入Hydra全局运行上下文，不需要修改插件源码即可直接获取。拿到trial实例后调用原生的set_user_attr方法，就能写入任意自定义信息。
直接修改目标函数代码即可，调整后的示例如下：

from hydra.core.hydra_config import HydraConfig
from hydra_plugins.hydra_optuna_sweeper.optuna_sweeper import _TrialCallback

@hydra.main(version_base=None, config_path="conf", config_name="sweep")
def sphere(cfg: TrainConfig) -> float:
    # 从当前运行上下文中提取当前绑定的trial实例
    current_trial = None
    for callback in HydraConfig.get().callbacks:
        if isinstance(callback, _TrialCallback):
            current_trial = callback.trial
            break

    # 替换为自身业务逻辑：生成自有数据库记录，拿到对应记录ID
    my_db_record_id = "your_custom_record_id"
    # 将自定义信息写入trial的用户属性
    if current_trial:
        current_trial.set_user_attr("my_db_record_id", my_db_record_id)

    x: float = cfg.x
    y: float = cfg.y
    return x**2 + y**2

注意事项 #

• 该方案适配Hydra Optuna Sweeper 1.2及以上版本，无侵入性，不需要修改插件本身的封装代码
• 写入的自定义属性后续可直接通过trial.user_attrs字段读取，无需再通过超参数组合反向匹配自有数据库记录
• 多进程并行调参（n_jobs>1）场景下该方法依然有效，每个执行进程只会拿到自身绑定的trial实例，不会出现数据错乱

内容的提问来源于stack exchange，提问作者Michel Kok