| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| DEFAULT_CONFIG = ROOT / "conf/config.yaml" |
|
|
|
|
| def load_config(path: str | Path = DEFAULT_CONFIG) -> dict[str, Any]: |
| config_path = Path(path).expanduser().resolve() |
| with config_path.open("r", encoding="utf-8") as handle: |
| config = yaml.safe_load(handle) |
| if not isinstance(config, dict): |
| raise ValueError("configuration must be a YAML mapping") |
| for section in ("data", "model", "train", "dataloader", "distributed"): |
| if section not in config: |
| raise ValueError(f"missing configuration section: {section}") |
| data, model = config["data"], config["model"] |
| for key in ("input_length", "output_length", "height", "width", "channels"): |
| if int(data.get(key, 0)) <= 0: |
| raise ValueError(f"data.{key} must be positive") |
| if int(data["height"]) % 4 or int(data["width"]) % 4: |
| raise ValueError("data height and width must be divisible by four") |
| dims, depths = model.get("dims"), model.get("depths") |
| if not isinstance(dims, list) or len(dims) != 2 or int(dims[1]) != 2 * int(dims[0]): |
| raise ValueError("model.dims must be [D, 2*D]") |
| if not isinstance(depths, list) or len(depths) != 2 or min(int(x) for x in depths) < 1: |
| raise ValueError("model.depths must contain two positive integers") |
| heads = int(model.get("heads", 0)) |
| if heads < 1 or any(int(dim) % heads for dim in dims): |
| raise ValueError("model.heads must divide both hidden dimensions") |
| normalization = data.get("normalization", "unit") |
| if normalization not in ("unit", "uint8_255"): |
| raise ValueError("data.normalization must be 'unit' or 'uint8_255'") |
| for key in ("data_dir", "train_npz", "val_npz", "test_npz"): |
| if data.get(key): |
| value = Path(data[key]).expanduser() |
| data[key] = str(value if value.is_absolute() else ROOT / value) |
| output_dir = Path(config["train"]["output_dir"]).expanduser() |
| config["train"]["output_dir"] = str(output_dir if output_dir.is_absolute() else ROOT / output_dir) |
| return config |
|
|
|
|
| def resolve_device(requested: str, local_rank: int = 0) -> torch.device: |
| if requested not in ("auto", "cpu", "cuda"): |
| raise ValueError("device must be auto, cpu, or cuda") |
| use_accelerator = requested == "cuda" or (requested == "auto" and torch.cuda.is_available()) |
| if use_accelerator: |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA/ROCm device requested but torch.cuda.is_available() is false") |
| torch.cuda.set_device(local_rank) |
| return torch.device("cuda", local_rank) |
| return torch.device("cpu") |
|
|
|
|
| def resolve_cli_path(value: str | None) -> str | None: |
| if not value: |
| return value |
| path = Path(value).expanduser() |
| return str(path if path.is_absolute() else ROOT / path) |
|
|
|
|
| def load_checkpoint_payload(path: str | Path, device: torch.device) -> dict[str, Any]: |
| payload = torch.load(path, map_location=device, weights_only=False) |
| if not isinstance(payload, dict) or "model" not in payload or "config" not in payload: |
| raise ValueError("checkpoint must contain model and config") |
| return payload |
|
|
|
|
| def clean_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| return {key.removeprefix("module."): value for key, value in state_dict.items()} |
|
|
|
|
| def atomic_torch_save(payload: dict[str, Any], path: str | Path) -> Path: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") |
| torch.save(payload, temporary) |
| os.replace(temporary, path) |
| return path |
|
|
