Upload landmarkdiff/hyperparam.py with huggingface_hub

landmarkdiff/hyperparam.py

@@ -0,0 +1,330 @@
+"""Hyperparameter search utilities for systematic ControlNet tuning.
+
+Supports grid search, random search, and Bayesian-inspired adaptive search
+over training hyperparameters. Generates YAML configs for each trial and
+tracks results for comparison.
+
+Usage:
+    from landmarkdiff.hyperparam import HyperparamSearch, SearchSpace
+
+    space = SearchSpace()
+    space.add_float("learning_rate", 1e-6, 1e-4, log_scale=True)
+    space.add_choice("optimizer", ["adamw", "adam8bit"])
+    space.add_int("batch_size", 2, 8, step=2)
+
+    search = HyperparamSearch(space, output_dir="hp_search")
+    for trial in search.generate_trials(strategy="random", n_trials=20):
+        print(trial.config)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Iterator
+
+
+def _to_native(val: Any) -> Any:
+    """Convert numpy/non-standard types to native Python for YAML serialization."""
+    if hasattr(val, "item"):  # numpy scalar
+        return val.item()
+    return val
+
+
+@dataclass
+class ParamSpec:
+    """Specification for a single hyperparameter."""
+
+    name: str
+    param_type: str  # "float", "int", "choice"
+    low: float | None = None
+    high: float | None = None
+    step: float | None = None
+    log_scale: bool = False
+    choices: list[Any] | None = None
+
+    def sample(self, rng) -> Any:
+        """Sample a value from this parameter spec."""
+        if self.param_type == "choice":
+            return rng.choice(self.choices)
+        elif self.param_type == "float":
+            if self.log_scale:
+                log_low = math.log(self.low)
+                log_high = math.log(self.high)
+                return float(math.exp(rng.uniform(log_low, log_high)))
+            return float(rng.uniform(self.low, self.high))
+        elif self.param_type == "int":
+            if self.step and self.step > 1:
+                n_steps = int((self.high - self.low) / self.step) + 1
+                idx = rng.integers(0, n_steps)
+                return int(self.low + idx * self.step)
+            return int(rng.integers(int(self.low), int(self.high) + 1))
+        raise ValueError(f"Unknown param type: {self.param_type}")
+
+    def grid_values(self, n_points: int = 5) -> list[Any]:
+        """Generate grid values for this parameter."""
+        if self.param_type == "choice":
+            return list(self.choices)
+        elif self.param_type == "int":
+            if self.step and self.step > 1:
+                vals = []
+                v = self.low
+                while v <= self.high:
+                    vals.append(int(v))
+                    v += self.step
+                return vals
+            return list(range(int(self.low), int(self.high) + 1))
+        elif self.param_type == "float":
+            if self.log_scale:
+                log_low = math.log(self.low)
+                log_high = math.log(self.high)
+                return [
+                    float(math.exp(log_low + i * (log_high - log_low) / (n_points - 1)))
+                    for i in range(n_points)
+                ]
+            return [
+                float(self.low + i * (self.high - self.low) / (n_points - 1))
+                for i in range(n_points)
+            ]
+        return []
+
+
+class SearchSpace:
+    """Define the hyperparameter search space."""
+
+    def __init__(self) -> None:
+        self.params: dict[str, ParamSpec] = {}
+
+    def add_float(
+        self, name: str, low: float, high: float, log_scale: bool = False,
+    ) -> SearchSpace:
+        """Add a continuous float parameter."""
+        self.params[name] = ParamSpec(
+            name=name, param_type="float", low=low, high=high, log_scale=log_scale,
+        )
+        return self
+
+    def add_int(
+        self, name: str, low: int, high: int, step: int = 1,
+    ) -> SearchSpace:
+        """Add an integer parameter."""
+        self.params[name] = ParamSpec(
+            name=name, param_type="int", low=low, high=high, step=step,
+        )
+        return self
+
+    def add_choice(self, name: str, choices: list[Any]) -> SearchSpace:
+        """Add a categorical parameter."""
+        self.params[name] = ParamSpec(
+            name=name, param_type="choice", choices=choices,
+        )
+        return self
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def __contains__(self, name: str) -> bool:
+        return name in self.params
+
+
+@dataclass
+class Trial:
+    """A single hyperparameter trial."""
+
+    trial_id: str
+    config: dict[str, Any]
+    result: dict[str, float] = field(default_factory=dict)
+    status: str = "pending"  # pending, running, completed, failed
+
+    @property
+    def config_hash(self) -> str:
+        """Short hash of the config for deduplication."""
+        s = json.dumps(self.config, sort_keys=True, default=str)
+        return hashlib.md5(s.encode()).hexdigest()[:8]
+
+
+class HyperparamSearch:
+    """Hyperparameter search engine.
+
+    Args:
+        space: Search space definition.
+        output_dir: Directory to save trial configs and results.
+        seed: Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        space: SearchSpace,
+        output_dir: str | Path = "hp_search",
+        seed: int = 42,
+    ) -> None:
+        self.space = space
+        self.output_dir = Path(output_dir)
+        self.seed = seed
+        self.trials: list[Trial] = []
+
+    def generate_trials(
+        self,
+        strategy: str = "random",
+        n_trials: int = 20,
+        grid_points: int = 5,
+    ) -> list[Trial]:
+        """Generate trial configurations.
+
+        Args:
+            strategy: "random" or "grid".
+            n_trials: Number of trials for random search.
+            grid_points: Points per continuous dimension for grid search.
+
+        Returns:
+            List of Trial objects with configs.
+        """
+        if strategy == "grid":
+            trials = self._grid_search(grid_points)
+        elif strategy == "random":
+            trials = self._random_search(n_trials)
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}. Use 'random' or 'grid'.")
+
+        self.trials.extend(trials)
+        return trials
+
+    def _random_search(self, n_trials: int) -> list[Trial]:
+        """Generate random trial configs."""
+        import numpy as np
+
+        rng = np.random.default_rng(self.seed)
+        seen_hashes: set[str] = set()
+        trials: list[Trial] = []
+
+        max_attempts = n_trials * 10
+        attempts = 0
+        while len(trials) < n_trials and attempts < max_attempts:
+            attempts += 1
+            config = {
+                name: spec.sample(rng)
+                for name, spec in self.space.params.items()
+            }
+            trial = Trial(
+                trial_id=f"trial_{len(trials):04d}",
+                config=config,
+            )
+            if trial.config_hash not in seen_hashes:
+                seen_hashes.add(trial.config_hash)
+                trials.append(trial)
+
+        return trials
+
+    def _grid_search(self, grid_points: int) -> list[Trial]:
+        """Generate grid search configs."""
+        import itertools
+
+        param_names = list(self.space.params.keys())
+        param_values = [
+            self.space.params[name].grid_values(grid_points)
+            for name in param_names
+        ]
+
+        trials = []
+        for combo in itertools.product(*param_values):
+            config = dict(zip(param_names, combo))
+            trial = Trial(
+                trial_id=f"trial_{len(trials):04d}",
+                config=config,
+            )
+            trials.append(trial)
+
+        return trials
+
+    def record_result(
+        self, trial_id: str, metrics: dict[str, float],
+    ) -> None:
+        """Record results for a trial."""
+        for trial in self.trials:
+            if trial.trial_id == trial_id:
+                trial.result = metrics
+                trial.status = "completed"
+                return
+        raise KeyError(f"Trial {trial_id} not found")
+
+    def best_trial(
+        self, metric: str = "loss", lower_is_better: bool = True,
+    ) -> Trial | None:
+        """Get the best completed trial by a metric."""
+        completed = [t for t in self.trials if t.status == "completed" and metric in t.result]
+        if not completed:
+            return None
+        return (min if lower_is_better else max)(completed, key=lambda t: t.result[metric])
+
+    def save_configs(self) -> Path:
+        """Save all trial configs as YAML files.
+
+        Returns:
+            Output directory path.
+        """
+        import yaml
+
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        for trial in self.trials:
+            cfg_path = self.output_dir / f"{trial.trial_id}.yaml"
+            # Convert numpy types to native Python for YAML serialization
+            native_config = {k: _to_native(v) for k, v in trial.config.items()}
+            with open(cfg_path, "w") as f:
+                yaml.safe_dump(
+                    {"trial_id": trial.trial_id, **native_config},
+                    f, default_flow_style=False,
+                )
+
+        # Save summary index
+        index = {
+            "seed": self.seed,
+            "n_trials": len(self.trials),
+            "params": {
+                name: {
+                    "type": spec.param_type,
+                    "low": spec.low,
+                    "high": spec.high,
+                    "choices": spec.choices,
+                    "log_scale": spec.log_scale,
+                }
+                for name, spec in self.space.params.items()
+            },
+        }
+        with open(self.output_dir / "search_index.json", "w") as f:
+            json.dump(index, f, indent=2, default=str)
+
+        return self.output_dir
+
+    def results_table(self) -> str:
+        """Format results as a text table."""
+        completed = [t for t in self.trials if t.status == "completed"]
+        if not completed:
+            return "No completed trials."
+
+        # Collect all metric names
+        metric_names = sorted(set().union(*(t.result.keys() for t in completed)))
+        param_names = sorted(self.space.params.keys())
+
+        # Header
+        cols = ["Trial"] + param_names + metric_names
+        lines = [" | ".join(f"{c:>12s}" for c in cols)]
+        lines.append("-" * len(lines[0]))
+
+        # Rows
+        for trial in completed:
+            parts = [f"{trial.trial_id:>12s}"]
+            for p in param_names:
+                val = trial.config.get(p, "")
+                if isinstance(val, float):
+                    parts.append(f"{val:>12.6f}")
+                else:
+                    parts.append(f"{str(val):>12s}")
+            for m in metric_names:
+                val = trial.result.get(m, float("nan"))
+                parts.append(f"{val:>12.4f}")
+            lines.append(" | ".join(parts))
+
+        return "\n".join(lines)