Spaces:
Runtime error
Runtime error
| import os | |
| from dataclasses import dataclass, field | |
| from datetime import datetime | |
| from omegaconf import OmegaConf | |
| import craftsman | |
| from craftsman.utils.typing import * | |
| # ============ Register OmegaConf Recolvers ============= # | |
| OmegaConf.register_new_resolver( | |
| "calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n) | |
| ) | |
| OmegaConf.register_new_resolver("add", lambda a, b: a + b) | |
| OmegaConf.register_new_resolver("sub", lambda a, b: a - b) | |
| OmegaConf.register_new_resolver("mul", lambda a, b: a * b) | |
| OmegaConf.register_new_resolver("div", lambda a, b: a / b) | |
| OmegaConf.register_new_resolver("idiv", lambda a, b: a // b) | |
| OmegaConf.register_new_resolver("basename", lambda p: os.path.basename(p)) | |
| OmegaConf.register_new_resolver("rmspace", lambda s, sub: str(s).replace(" ", sub)) | |
| OmegaConf.register_new_resolver("tuple2", lambda s: [float(s), float(s)]) | |
| OmegaConf.register_new_resolver("gt0", lambda s: s > 0) | |
| OmegaConf.register_new_resolver("cmaxgt0", lambda s: C_max(s) > 0) | |
| OmegaConf.register_new_resolver("not", lambda s: not s) | |
| OmegaConf.register_new_resolver( | |
| "cmaxgt0orcmaxgt0", lambda a, b: C_max(a) > 0 or C_max(b) > 0 | |
| ) | |
| # ======================================================= # | |
| def C_max(value: Any) -> float: | |
| if isinstance(value, int) or isinstance(value, float): | |
| pass | |
| else: | |
| value = config_to_primitive(value) | |
| if not isinstance(value, list): | |
| raise TypeError("Scalar specification only supports list, got", type(value)) | |
| if len(value) >= 6: | |
| max_value = value[2] | |
| for i in range(4, len(value), 2): | |
| max_value = max(max_value, value[i]) | |
| value = [value[0], value[1], max_value, value[3]] | |
| if len(value) == 3: | |
| value = [0] + value | |
| assert len(value) == 4 | |
| start_step, start_value, end_value, end_step = value | |
| value = max(start_value, end_value) | |
| return value | |
| class ExperimentConfig: | |
| name: str = "default" | |
| description: str = "" | |
| tag: str = "" | |
| seed: int = 0 | |
| use_timestamp: bool = True | |
| timestamp: Optional[str] = None | |
| exp_root_dir: str = "outputs" | |
| ### these shouldn't be set manually | |
| exp_dir: str = "outputs/default" | |
| trial_name: str = "exp" | |
| trial_dir: str = "outputs/default/exp" | |
| n_gpus: int = 1 | |
| ### | |
| resume: Optional[str] = None | |
| data_type: str = "" | |
| data: dict = field(default_factory=dict) | |
| system_type: str = "" | |
| system: dict = field(default_factory=dict) | |
| # accept pytorch-lightning trainer parameters | |
| # see https://lightning.ai/docs/pytorch/stable/common/trainer.html#trainer-class-api | |
| trainer: dict = field(default_factory=dict) | |
| # accept pytorch-lightning checkpoint callback parameters | |
| # see https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint | |
| checkpoint: dict = field(default_factory=dict) | |
| def __post_init__(self): | |
| if not self.tag and not self.use_timestamp: | |
| raise ValueError("Either tag is specified or use_timestamp is True.") | |
| self.trial_name = self.tag | |
| # if resume from an existing config, self.timestamp should not be None | |
| if self.timestamp is None: | |
| self.timestamp = "" | |
| if self.use_timestamp: | |
| if self.n_gpus > 1: | |
| craftsman.warn( | |
| "Timestamp is disabled when using multiple GPUs, please make sure you have a unique tag." | |
| ) | |
| else: | |
| self.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S") | |
| self.trial_name += self.timestamp | |
| self.exp_dir = os.path.join(self.exp_root_dir, self.name) | |
| self.trial_dir = os.path.join(self.exp_dir, self.trial_name) | |
| os.makedirs(self.trial_dir, exist_ok=True) | |
| def load_config(*yamls: str, cli_args: list = [], from_string=False, **kwargs) -> Any: | |
| if from_string: | |
| yaml_confs = [OmegaConf.create(s) for s in yamls] | |
| else: | |
| yaml_confs = [OmegaConf.load(f) for f in yamls] | |
| cli_conf = OmegaConf.from_cli(cli_args) | |
| cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs) | |
| OmegaConf.resolve(cfg) | |
| assert isinstance(cfg, DictConfig) | |
| scfg = parse_structured(ExperimentConfig, cfg) | |
| return scfg | |
| def config_to_primitive(config, resolve: bool = True) -> Any: | |
| return OmegaConf.to_container(config, resolve=resolve) | |
| def dump_config(path: str, config) -> None: | |
| with open(path, "w") as fp: | |
| OmegaConf.save(config=config, f=fp) | |
| def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any: | |
| scfg = OmegaConf.structured(fields(**cfg)) | |
| return scfg |