342 lines
11 KiB
Python
342 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import tomllib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunConfig:
|
|
name: str
|
|
seed: int
|
|
artifact_dir: Path
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DataConfig:
|
|
root: Path
|
|
train_cases: int
|
|
val_cases: int
|
|
test_cases: int
|
|
points_per_case: int
|
|
batch_size: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ModelConfig:
|
|
type: str
|
|
hidden_width: int
|
|
depth: int
|
|
activation: str
|
|
coordinate_features: tuple[str, ...]
|
|
fourier_scales: tuple[float, ...]
|
|
condition_width: int
|
|
condition_depth: int
|
|
condition_dim: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OptimConfig:
|
|
lr: float
|
|
weight_decay: float
|
|
steps: int
|
|
log_interval: int | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CheckpointConfig:
|
|
interval_seconds: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class StabilityConfig:
|
|
max_grad_norm: float | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PrecisionConfig:
|
|
dtype: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DeviceConfig:
|
|
type: str
|
|
allow_cpu_fallback: bool
|
|
benchmark_kernels: bool
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LossConfig:
|
|
type: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ObservabilityConfig:
|
|
backend: str
|
|
project: str
|
|
entity: str | None
|
|
mode: str
|
|
tags: tuple[str, ...]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TrainingConfig:
|
|
path: Path
|
|
config_text: str
|
|
run: RunConfig
|
|
data: DataConfig
|
|
model: ModelConfig
|
|
optim: OptimConfig
|
|
device: DeviceConfig
|
|
loss: LossConfig
|
|
checkpoint: CheckpointConfig
|
|
stability: StabilityConfig
|
|
precision: PrecisionConfig
|
|
observability: ObservabilityConfig
|
|
|
|
_REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss")
|
|
|
|
|
|
def load_training_config(path: str | Path) -> TrainingConfig:
|
|
config_path = Path(path).expanduser()
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"Training config not found: {config_path}")
|
|
if not config_path.is_file():
|
|
raise ValueError(f"Training config is not a file: {config_path}")
|
|
|
|
text = config_path.read_text()
|
|
try:
|
|
raw = tomllib.loads(text)
|
|
except tomllib.TOMLDecodeError as exc:
|
|
raise ValueError(f"Invalid TOML config {config_path}: {exc}") from exc
|
|
|
|
for section_name in _REQUIRED_SECTIONS:
|
|
if section_name not in raw or not isinstance(raw[section_name], dict):
|
|
raise ValueError(f"Training config missing [{section_name}] section")
|
|
|
|
run_raw = raw["run"]
|
|
data_raw = raw["data"]
|
|
model_raw = raw["model"]
|
|
optim_raw = raw["optim"]
|
|
device_raw = raw["device"]
|
|
loss_raw = raw["loss"]
|
|
checkpoint_raw = raw.get("checkpoint", {})
|
|
if checkpoint_raw is None:
|
|
checkpoint_raw = {}
|
|
if not isinstance(checkpoint_raw, dict):
|
|
raise ValueError("Training config [checkpoint] section must be a table")
|
|
stability_raw = raw.get("stability", {})
|
|
if stability_raw is None:
|
|
stability_raw = {}
|
|
if not isinstance(stability_raw, dict):
|
|
raise ValueError("Training config [stability] section must be a table")
|
|
precision_raw = raw.get("precision", {})
|
|
if precision_raw is None:
|
|
precision_raw = {}
|
|
if not isinstance(precision_raw, dict):
|
|
raise ValueError("Training config [precision] section must be a table")
|
|
observability_raw = raw.get("observability", {})
|
|
if observability_raw is None:
|
|
observability_raw = {}
|
|
if not isinstance(observability_raw, dict):
|
|
raise ValueError("Training config [observability] section must be a table")
|
|
|
|
|
|
run = RunConfig(
|
|
name=_string(run_raw, "name"),
|
|
seed=_integer(run_raw, "seed", minimum=0),
|
|
artifact_dir=_path(run_raw, "artifact_dir"),
|
|
)
|
|
data = DataConfig(
|
|
root=_path(data_raw, "root"),
|
|
train_cases=_integer(data_raw, "train_cases", minimum=1),
|
|
val_cases=_integer(data_raw, "val_cases", minimum=0),
|
|
test_cases=_integer(data_raw, "test_cases", minimum=0),
|
|
points_per_case=_integer(data_raw, "points_per_case", minimum=1),
|
|
batch_size=_integer(data_raw, "batch_size", minimum=1),
|
|
)
|
|
model = ModelConfig(
|
|
type=_choice(_string(model_raw, "type"), {"mlp", "film_fourier_mlp"}, "model.type"),
|
|
hidden_width=_integer(model_raw, "hidden_width", minimum=1),
|
|
depth=_integer(model_raw, "depth", minimum=1),
|
|
activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"),
|
|
coordinate_features=_string_tuple(model_raw, "coordinate_features", default=("x", "y", "sdf")),
|
|
fourier_scales=_number_tuple(model_raw, "fourier_scales", default=(1.0, 2.0, 4.0, 8.0, 16.0)),
|
|
condition_width=_integer(model_raw, "condition_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
|
|
condition_depth=_integer(model_raw, "condition_depth", minimum=1, default=2),
|
|
condition_dim=_integer(model_raw, "condition_dim", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
|
|
)
|
|
optim = OptimConfig(
|
|
lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True),
|
|
weight_decay=_number(optim_raw, "weight_decay", minimum=0.0),
|
|
steps=_integer(optim_raw, "steps", minimum=1),
|
|
log_interval=_optional_integer(optim_raw, "log_interval", minimum=1),
|
|
)
|
|
device = DeviceConfig(
|
|
type=_choice(_string(device_raw, "type").lower(), {"cuda", "cpu", "auto"}, "device.type"),
|
|
allow_cpu_fallback=_boolean(device_raw, "allow_cpu_fallback"),
|
|
benchmark_kernels=_boolean(device_raw, "benchmark_kernels"),
|
|
)
|
|
loss = LossConfig(type=_choice(_string(loss_raw, "type"), {"normalized_mse"}, "loss.type"))
|
|
checkpoint = CheckpointConfig(
|
|
interval_seconds=_integer(checkpoint_raw, "interval_seconds", minimum=0, default=1800),
|
|
)
|
|
stability = StabilityConfig(
|
|
max_grad_norm=_optional_number(stability_raw, "max_grad_norm", minimum=0.0, exclusive_minimum=True),
|
|
)
|
|
precision = PrecisionConfig(
|
|
dtype=_choice(_string(precision_raw, "dtype", default="float32").lower(), {"float32", "bf16"}, "precision.dtype"),
|
|
)
|
|
observability = ObservabilityConfig(
|
|
backend=_choice(_string(observability_raw, "backend", default="none").lower(), {"none", "wandb"}, "observability.backend"),
|
|
project=_string(observability_raw, "project", default="airfrans"),
|
|
entity=_optional_string(observability_raw, "entity"),
|
|
mode=_choice(_string(observability_raw, "mode", default="online").lower(), {"online", "offline", "disabled"}, "observability.mode"),
|
|
tags=_string_tuple(observability_raw, "tags", default=()),
|
|
)
|
|
|
|
|
|
requested_cases = data.train_cases + data.val_cases + data.test_cases
|
|
if requested_cases <= 0:
|
|
raise ValueError("Training config must request at least one case")
|
|
|
|
return TrainingConfig(
|
|
path=config_path,
|
|
config_text=text,
|
|
run=run,
|
|
data=data,
|
|
model=model,
|
|
optim=optim,
|
|
device=device,
|
|
loss=loss,
|
|
checkpoint=checkpoint,
|
|
stability=stability,
|
|
precision=precision,
|
|
observability=observability,
|
|
)
|
|
|
|
|
|
def _path(section: dict[str, Any], key: str) -> Path:
|
|
value = _string(section, key)
|
|
path = Path(value).expanduser()
|
|
if path.is_absolute():
|
|
return path
|
|
return Path.cwd() / path
|
|
|
|
|
|
def _string(section: dict[str, Any], key: str, *, default: str | None = None) -> str:
|
|
if key not in section:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Training config missing key: {key}")
|
|
value = section[key]
|
|
if not isinstance(value, str) or not value:
|
|
raise ValueError(f"Expected non-empty string for {key}")
|
|
return value
|
|
|
|
|
|
def _optional_string(section: dict[str, Any], key: str) -> str | None:
|
|
if key not in section:
|
|
return None
|
|
return _string(section, key)
|
|
|
|
|
|
def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None, default: int | None = None) -> int:
|
|
if key not in section:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Training config missing key: {key}")
|
|
value = section[key]
|
|
if isinstance(value, bool) or not isinstance(value, int):
|
|
raise ValueError(f"Expected integer for {key}")
|
|
if minimum is not None and value < minimum:
|
|
raise ValueError(f"Expected {key} >= {minimum}")
|
|
return value
|
|
|
|
|
|
def _optional_integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int | None:
|
|
if key not in section:
|
|
return None
|
|
return _integer(section, key, minimum=minimum)
|
|
|
|
|
|
def _number(
|
|
section: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
minimum: float | None = None,
|
|
exclusive_minimum: bool = False,
|
|
) -> float:
|
|
value = _required(section, key)
|
|
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
|
raise ValueError(f"Expected number for {key}")
|
|
result = float(value)
|
|
if minimum is not None:
|
|
if exclusive_minimum and result <= minimum:
|
|
raise ValueError(f"Expected {key} > {minimum}")
|
|
if not exclusive_minimum and result < minimum:
|
|
raise ValueError(f"Expected {key} >= {minimum}")
|
|
return result
|
|
|
|
|
|
def _optional_number(
|
|
section: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
minimum: float | None = None,
|
|
exclusive_minimum: bool = False,
|
|
) -> float | None:
|
|
if key not in section:
|
|
return None
|
|
return _number(section, key, minimum=minimum, exclusive_minimum=exclusive_minimum)
|
|
|
|
|
|
def _number_tuple(section: dict[str, Any], key: str, *, default: tuple[float, ...]) -> tuple[float, ...]:
|
|
if key not in section:
|
|
return default
|
|
value = section[key]
|
|
if not isinstance(value, list) or not value:
|
|
raise ValueError(f"Expected non-empty number array for {key}")
|
|
result: list[float] = []
|
|
for item in value:
|
|
if isinstance(item, bool) or not isinstance(item, (int, float)):
|
|
raise ValueError(f"Expected number array for {key}")
|
|
result.append(float(item))
|
|
return tuple(result)
|
|
|
|
|
|
def _string_tuple(section: dict[str, Any], key: str, *, default: tuple[str, ...]) -> tuple[str, ...]:
|
|
if key not in section:
|
|
return default
|
|
value = section[key]
|
|
if not isinstance(value, list) or not value:
|
|
raise ValueError(f"Expected non-empty string array for {key}")
|
|
result: list[str] = []
|
|
for item in value:
|
|
if not isinstance(item, str) or not item:
|
|
raise ValueError(f"Expected non-empty string array for {key}")
|
|
result.append(item)
|
|
if len(set(result)) != len(result):
|
|
raise ValueError(f"Expected unique strings for {key}")
|
|
return tuple(result)
|
|
|
|
|
|
def _boolean(section: dict[str, Any], key: str) -> bool:
|
|
value = _required(section, key)
|
|
if not isinstance(value, bool):
|
|
raise ValueError(f"Expected boolean for {key}")
|
|
return value
|
|
|
|
|
|
def _choice(value: str, allowed: set[str], key: str) -> str:
|
|
if value not in allowed:
|
|
allowed_text = ", ".join(sorted(allowed))
|
|
raise ValueError(f"Expected {key} to be one of: {allowed_text}")
|
|
return value
|
|
|
|
|
|
def _required(section: dict[str, Any], key: str) -> Any:
|
|
if key not in section:
|
|
raise ValueError(f"Training config missing key: {key}")
|
|
return section[key]
|