File size: 3,226 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | from typing import Optional
import os
from omegaconf import OmegaConf
def _dict_to_str(d: dict) -> str:
"""
Convert a dictionary to a string without quotes.
"""
output = "{"
for key, value in d.items():
if value is None:
value = "null"
output += (
f"{key}: {_dict_to_str(value) if isinstance(value, dict) else value}, "
)
output = output[:-2] + "}"
return output
def _yaml_to_cli(
yaml_path: str,
prefix: Optional[str] = None,
) -> list[str]:
"""
Convert a yaml file to a list of command line arguments.
"""
cfg = OmegaConf.load(yaml_path)
cli = []
for key, value in OmegaConf.to_container(cfg).items():
if value is None:
value = "null"
cli.append(
f"++{prefix + '.' if prefix else ''}{key}={_dict_to_str(value) if isinstance(value, dict) else value}"
)
return cli
def unwrap_shortcuts(
argv: list[str],
config_path: str,
config_name: str,
) -> list[str]:
"""
Unwrap shortcuts by replacing them with commands from corresponding yaml files.
All shortcuts should be in the form of `@shortcut_name`.
Example:
- @latent -> unwrap configurations/shortcut/latent/base.yaml and configurations/shortcut/latent/dataset_name.yaml
- @mit_vision/h100 -> unwrap configurations/shortcut/mit_vision/h100.yaml
"""
# find the default dataset
defaults = OmegaConf.load(f"{config_path}/{config_name}.yaml").defaults
dataset = next(default.dataset for default in defaults if "dataset" in default)
# check if dataset is overridden
for arg in argv:
if arg.startswith("dataset="):
dataset = arg.split("=")[1]
if dataset is None:
raise ValueError("Dataset name is not provided.")
new_argv = []
for arg in argv:
if arg.startswith("@"):
shortcut = arg[1:]
base_path = f"{config_path}/shortcut/{shortcut}/base.yaml"
if os.path.exists(base_path):
new_argv += _yaml_to_cli(base_path)
dataset_path = f"{config_path}/shortcut/{shortcut}/{dataset}.yaml"
if os.path.exists(dataset_path):
new_argv += _yaml_to_cli(dataset_path)
else:
default_path = f"{config_path}/shortcut/{shortcut}.yaml"
if os.path.exists(default_path):
new_argv += _yaml_to_cli(default_path)
else:
raise ValueError(f"Shortcut @{shortcut} not found.")
elif arg.startswith("algorithm/backbone="):
# this is a workaround to enable overriding the backbone in the command line
# otherwise, the backbone could be re-overridden by
# the backbone cfgs in dataset-experiment dependent cfgs
new_argv += override_backbone(arg[19:])
else:
new_argv.append(arg)
return new_argv
def override_backbone(name: str) -> list[str]:
"""
Override the backbone with the specified name.
"""
return ["algorithm.backbone=null"] + _yaml_to_cli(
f"configurations/algorithm/backbone/{name}.yaml", prefix="algorithm.backbone"
)
|