Files
optuna-dashboard/optuna_dashboard/serializer.py
T
2020-10-27 03:12:35 +09:00

111 lines
3.1 KiB
Python

import json
from typing import Any, Dict, List
from optuna.study import StudySummary
from optuna.trial import FrozenTrial
try:
from typing import TypedDict
except:
from typing_extensions import TypedDict
Attribute = TypedDict(
"Attribute",
{
"key": str,
"value": str,
},
)
IntermediateValue = TypedDict(
"IntermediateValue",
{
"step": int,
"value": float,
},
)
TrialParam = TypedDict(
"TrialParam",
{
"name": str,
"value": str,
},
)
def serialize_attrs(attrs: Dict[str, Any]) -> List[Attribute]:
return [{"key": k, "value": json.dumps(v)} for k, v in attrs.items()]
def serialize_intermediate_values(values: Dict[int, float]) -> List[IntermediateValue]:
return [{"step": step, "value": value} for step, value in values.items()]
def serialize_trial_params(params: Dict[str, Any]) -> List[TrialParam]:
return [{"name": name, "value": str(value)} for name, value in params.items()]
def serialize_study_summary(summary: StudySummary) -> Dict[str, Any]:
serialized = {
"study_id": summary._study_id,
"study_name": summary.study_name,
"direction": summary.direction.name.lower(),
"user_attrs": serialize_attrs(summary.user_attrs),
"system_attrs": serialize_attrs(summary.system_attrs),
}
if summary.datetime_start is not None:
serialized["datetime_start"] = (summary.datetime_start.isoformat(),)
if summary.best_trial:
serialized["best_trial"] = serialize_frozen_trial(
summary._study_id, summary.best_trial
)
return serialized
def serialize_study_detail(
summary: StudySummary, trials: List[FrozenTrial]
) -> Dict[str, Any]:
serialized: Dict[str, Any] = {
"name": summary.study_name,
"direction": summary.direction.name.lower(),
}
if summary.datetime_start is not None:
serialized["datetime_start"] = summary.datetime_start.isoformat()
if summary.best_trial is not None:
serialized["best_trial"] = serialize_frozen_trial(
summary._study_id, summary.best_trial
)
serialized["trials"] = [
serialize_frozen_trial(summary._study_id, trial) for trial in trials
]
return serialized
def serialize_frozen_trial(study_id: int, trial: FrozenTrial) -> Dict[str, Any]:
serialized = {
"trial_id": trial._trial_id,
"study_id": study_id,
"number": trial.number,
"state": trial.state.name.capitalize(),
"intermediate_values": serialize_intermediate_values(trial.intermediate_values),
"params": serialize_trial_params(trial.params),
"user_attrs": serialize_attrs(trial.user_attrs),
"system_attrs": serialize_attrs(trial.system_attrs),
}
if trial.value is not None:
serialized["value"] = trial.value
if trial.datetime_start is not None:
serialized["datetime_start"] = trial.datetime_start.isoformat()
if trial.datetime_complete is not None:
serialized["datetime_complete"] = trial.datetime_complete.isoformat()
return serialized