diff --git a/configs/training/cifar10_default.yaml b/configs/training/cifar10_default.yaml index 1507486..a9d292a 100644 --- a/configs/training/cifar10_default.yaml +++ b/configs/training/cifar10_default.yaml @@ -1,7 +1,7 @@ # @package _global_ training: - num_epochs: 100 # CIFAR-10 needs more epochs than MNIST + num_epochs: 3 # CIFAR-10 needs more epochs than MNIST step_size: 0.05 # ODE solver step size (20 steps for inference) n_samples: 4 # Generate multiple samples at once log_interval: 10 # Log every 10 epochs diff --git a/configs/training/cifar10_improved.yaml b/configs/training/cifar10_improved.yaml index b84ddba..9e7c9ba 100644 --- a/configs/training/cifar10_improved.yaml +++ b/configs/training/cifar10_improved.yaml @@ -1,7 +1,7 @@ # @package _global_ training: - num_epochs: 300 # 3x longer than default (100 epochs) + num_epochs: 3 # 3x longer than default (100 epochs) step_size: 0.01 # 100 ODE steps for high-quality inference (vs 20 steps) n_samples: 9 # Generate 3x3 grid of samples log_interval: 10 # Log every 10 epochs diff --git a/experiments/_common.py b/experiments/_common.py index 35d1ef7..5694b3d 100644 --- a/experiments/_common.py +++ b/experiments/_common.py @@ -3,6 +3,7 @@ import matplotlib.pyplot as plt import numpy as np from tinygrad.nn.state import safe_load +from tinygrad.tensor import Tensor as T def detect_time_embed_dim(model_path: str, in_channels: int) -> int: @@ -43,3 +44,59 @@ def save_grid_figure(x_np: np.ndarray, title: str, path: str, figsize=(4, 4.2)): plt.tight_layout() plt.savefig(path, dpi=150, bbox_inches="tight") plt.close(fig) + + +def per_sample_norm_mean(x: T) -> float: + """Mean over the batch of the per-sample L2 norm.""" + flat = x.reshape(x.shape[0], -1) + return ((flat * flat).sum(axis=-1) + 1e-20).sqrt().mean().numpy().item() + + +def normalize_per_sample(x: T) -> T: + """Normalizes each sample in the batch to unit L2 norm.""" + flat = x.reshape(x.shape[0], -1) + mag = ((flat * flat).sum(axis=-1) + 1e-20).sqrt() + view_shape = (x.shape[0],) + (1,) * (len(x.shape) - 1) + return x / mag.reshape(view_shape) + + +def compute_fid_vs_real( + generated_np: np.ndarray, dataset_name: str, n_samples: int, weights_dir: str = "weights" +) -> float | None: + """FID of generated images against real images freshly sampled from `dataset_name`.""" + from tinyflow.dataloader import CIFAR10Loader, FashionMNISTLoader, MNISTLoader + from tinyflow.metrics import calculate_fid, get_feature_extractor + + classifier = get_feature_extractor(dataset_name, weights_dir=weights_dir) + if not classifier._weights_loaded: + print(f"Warning: FID classifier weights not found for {dataset_name}, skipping FID.") + return None + + if dataset_name == "mnist": + dataloader = MNISTLoader( + path="dataset/mnist/trainset/trainingSet/*/*.jpg", + batch_size=64, + shuffle=True, + flatten=False, + ) + elif dataset_name == "fashion_mnist": + dataloader = FashionMNISTLoader( + path="dataset/fashion_mnist", batch_size=64, shuffle=True, flatten=False, train=True + ) + elif dataset_name == "cifar10": + dataloader = CIFAR10Loader( + path="dataset/cifar10/cifar-10-batches-py", batch_size=64, shuffle=True, train=True + ) + else: + raise ValueError(f"Unknown dataset: {dataset_name}") + + real_images = [] + collected = 0 + for batch_images, _ in dataloader: + real_images.append(batch_images) + collected += len(batch_images) + if collected >= n_samples: + break + real_images = np.concatenate(real_images, axis=0)[:n_samples] + + return calculate_fid(real_images, generated_np, feature_extractor=classifier, batch_size=64) diff --git a/experiments/adaptive_sampler.py b/experiments/adaptive_sampler.py new file mode 100644 index 0000000..cec10a6 --- /dev/null +++ b/experiments/adaptive_sampler.py @@ -0,0 +1,250 @@ +""" +Experiment 1 (HOMOTOPY_EXPERIMENTS.md): Condition-Number-Adaptive Sampler. + +Replaces fixed-step Euler integration with curvature-aware step sizing: + + kappa ~= ||v_theta(x + eps*u, t) - v_theta(x, t)|| / eps (u: random unit dir) + dt <- dt_base / (1 + beta * kappa) + +Compares, at a matched NFE budget, three configs on MNIST: + reference fixed uniform Euler at --reference-nfe (gold standard, e.g. 50) + fixed fixed uniform Euler at the tested budget (control) + adaptive curvature-adaptive Euler at the tested budget + +Usage: + uv run experiments/adaptive_sampler.py + uv run experiments/adaptive_sampler.py --nfe-budgets 10 20 --beta 5.0 +""" + +import argparse +import os + +import matplotlib.pyplot as plt +import mlflow +import numpy as np +from _common import ( + compute_fid_vs_real, + detect_time_embed_dim, + normalize_per_sample, + per_sample_norm_mean, + save_grid_figure, + schedule_grid, +) +from tinygrad.tensor import Tensor as T + +from tinyflow.nn import UNetTinygrad +from tinyflow.trainer import BaseTrainer +from tinyflow.utils import preprocess_time_mnist + +plt.style.use("ggplot") +mlflow.set_tracking_uri("sqlite:///mlflow.db") + + +def integrate_fixed(model, preprocess, x0, N: int): + x = x0 + ts = schedule_grid(N, 1.0) + nfe = 0 + for i in range(len(ts) - 1): + t = (T.zeros(1) + float(ts[i])).contiguous() + t_in = preprocess(t, x) + v = model(x, t_in).realize() + nfe += 1 + h = float(ts[i + 1] - ts[i]) + x = (x + h * v).realize() + return x, {"steps": N, "nfe": nfe} + + +def integrate_adaptive( + model, + preprocess, + x0, + nfe_target: int, + beta: float, + eps_x: float, + kappa_every: int, + max_steps: int = 500, +): + """Runs to completion (t reaches 1), unlike a fixed-N loop. `nfe_target` + only sets the nominal dt_base (1/nfe_target); curvature can make the + *actual* NFE spent come out higher or lower than that nominal target — + that actual count is what should be compared against a fixed baseline.""" + x = x0 + t_val = 0.0 + nfe = 0 + steps = 0 + kappa = 0.0 + dt_base = 1.0 / nfe_target + log = {"t": [], "dt": [], "kappa": []} + + while t_val < 1.0 - 1e-9 and steps < max_steps: + t = (T.zeros(1) + t_val).contiguous() + t_in = preprocess(t, x) + v = model(x, t_in).realize() + nfe += 1 + + if steps % kappa_every == 0: + u_hat = normalize_per_sample(T.randn(*x.shape).realize()) + v_pert = model(x + eps_x * u_hat, t_in).realize() + nfe += 1 + kappa = per_sample_norm_mean((v_pert - v) * (1.0 / eps_x)) + + dt = dt_base / (1.0 + beta * kappa) + dt = min(dt, 1.0 - t_val) + x = (x + dt * v).realize() + t_val += dt + steps += 1 + log["t"].append(t_val) + log["dt"].append(dt) + log["kappa"].append(kappa) + + return x, {"steps": steps, "nfe": nfe, **log} + + +def make_dt_plot(log: dict, budget: int) -> plt.Figure: + fig, axes = plt.subplots(1, 2, figsize=(11, 4)) + axes[0].plot(log["t"], log["dt"], marker="o", markersize=3) + axes[0].set_xlabel("t") + axes[0].set_ylabel("dt") + axes[0].set_title(f"Adaptive step size (budget NFE={budget})") + + axes[1].plot(log["t"], log["kappa"], marker="o", markersize=3, color="C1") + axes[1].set_xlabel("t") + axes[1].set_ylabel("kappa (curvature estimate)") + axes[1].set_title("Curvature along trajectory") + + plt.tight_layout() + return fig + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", default="weights/model_mnist_unet_linear.safetensors") + parser.add_argument( + "--nfe-budgets", + type=int, + nargs="+", + default=[10, 20], + help="Nominal NFE targets (sets dt_base=1/target); adaptive runs to " + "completion so actual NFE spent may differ, and the fixed baseline " + "is matched to that actual value post-hoc.", + ) + parser.add_argument("--reference-nfe", type=int, default=50) + parser.add_argument("--beta", type=float, default=5.0) + parser.add_argument("--kappa-every", type=int, default=4) + parser.add_argument("--eps-x", type=float, default=1e-2) + parser.add_argument("--n-samples", type=int, default=256) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--experiment", default="homotopy_experiments") + args = parser.parse_args() + + if not os.path.exists(args.model_path): + raise FileNotFoundError(f"Model file not found: {args.model_path}") + + in_ch = 1 + time_embed_dim = detect_time_embed_dim(args.model_path, in_ch) + model = UNetTinygrad(in_ch, in_ch, time_embed_dim=time_embed_dim) + BaseTrainer.load_model(model, args.model_path) + T.training = False + + mlflow.set_experiment(args.experiment) + with mlflow.start_run(run_name="1_adaptive_sampler") as parent: + mlflow.set_tags({"exp": "1_adaptive_sampler", "dataset": "mnist"}) + mlflow.log_params( + { + "model_path": args.model_path, + "reference_nfe": args.reference_nfe, + "beta": args.beta, + "kappa_every": args.kappa_every, + "eps_x": args.eps_x, + "n_samples": args.n_samples, + "seed": args.seed, + "nfe_budgets": ",".join(str(b) for b in args.nfe_budgets), + } + ) + + # gold-standard reference + T.manual_seed(args.seed) + x0_ref = T.randn(args.n_samples, 1, 28, 28).realize() + with mlflow.start_run(run_name=f"reference_NFE{args.reference_nfe}", nested=True): + x_ref, info = integrate_fixed(model, preprocess_time_mnist, x0_ref, args.reference_nfe) + x_ref_np = x_ref.numpy() + fid_ref = compute_fid_vs_real(x_ref_np, "mnist", args.n_samples) + mlflow.log_params({"config": "reference", **info}) + if fid_ref is not None: + mlflow.log_metric("fid", fid_ref) + grid_path = "/tmp/adaptive_sampler_reference.png" + save_grid_figure(x_ref_np, f"Reference NFE={args.reference_nfe}", grid_path) + mlflow.log_artifact(grid_path) + print(f"reference NFE={args.reference_nfe}: fid={fid_ref}") + + for target in args.nfe_budgets: + # adaptive runs to completion; the target only sets the nominal + # dt_base, actual NFE spent is discovered here + T.manual_seed(args.seed) + x0 = T.randn(args.n_samples, 1, 28, 28).realize() + + with mlflow.start_run(run_name=f"adaptive_target{target}", nested=True): + x_adapt, info = integrate_adaptive( + model, + preprocess_time_mnist, + x0, + target, + args.beta, + args.eps_x, + args.kappa_every, + ) + actual_nfe = info["nfe"] + x_adapt_np = x_adapt.numpy() + fid_adapt = compute_fid_vs_real(x_adapt_np, "mnist", args.n_samples) + mlflow.log_params( + { + "config": "adaptive", + "nfe_target": target, + "steps": info["steps"], + "nfe": actual_nfe, + "beta": args.beta, + "kappa_every": args.kappa_every, + } + ) + if fid_adapt is not None: + mlflow.log_metric("fid", fid_adapt) + mlflow.log_metric("mean_dt", float(np.mean(info["dt"]))) + mlflow.log_metric("mean_kappa", float(np.mean(info["kappa"]))) + + grid_path = f"/tmp/adaptive_sampler_adaptive_target{target}.png" + save_grid_figure( + x_adapt_np, f"Adaptive target={target} (actual NFE={actual_nfe})", grid_path + ) + mlflow.log_artifact(grid_path) + + dt_fig = make_dt_plot(info, target) + mlflow.log_figure(dt_fig, f"dt_profile_target{target}.png") + plt.close(dt_fig) + + print( + f"adaptive target={target} -> actual steps={info['steps']}, " + f"actual nfe={actual_nfe}: fid={fid_adapt}" + ) + + # fixed baseline at the SAME actual NFE the adaptive run spent + T.manual_seed(args.seed) + x0 = T.randn(args.n_samples, 1, 28, 28).realize() + + with mlflow.start_run(run_name=f"fixed_NFE{actual_nfe}", nested=True): + x_fixed, info = integrate_fixed(model, preprocess_time_mnist, x0, actual_nfe) + x_fixed_np = x_fixed.numpy() + fid_fixed = compute_fid_vs_real(x_fixed_np, "mnist", args.n_samples) + mlflow.log_params({"config": "fixed", "matched_to_target": target, **info}) + if fid_fixed is not None: + mlflow.log_metric("fid", fid_fixed) + grid_path = f"/tmp/adaptive_sampler_fixed_NFE{actual_nfe}.png" + save_grid_figure(x_fixed_np, f"Fixed NFE={actual_nfe}", grid_path) + mlflow.log_artifact(grid_path) + print(f"fixed NFE={actual_nfe} (matched to target={target}): fid={fid_fixed}") + + print(f"parent run: {parent.info.run_id}") + print("view with: uv run mlflow ui (tracking uri: sqlite:///mlflow.db)") + + +if __name__ == "__main__": + main() diff --git a/experiments/curvature_profile.py b/experiments/curvature_profile.py index f87f994..cca2ca8 100644 --- a/experiments/curvature_profile.py +++ b/experiments/curvature_profile.py @@ -19,7 +19,12 @@ import matplotlib.pyplot as plt import mlflow import numpy as np -from _common import detect_time_embed_dim, schedule_grid +from _common import ( + detect_time_embed_dim, + normalize_per_sample, + per_sample_norm_mean, + schedule_grid, +) from tinygrad.tensor import Tensor as T from tqdm import tqdm @@ -75,18 +80,6 @@ def build_model(dataset: str, model_path: str): return UNetTinygrad(in_ch, in_ch, time_embed_dim=detect_time_embed_dim(model_path, in_ch)) -def per_sample_norm_mean(x: T) -> float: - flat = x.reshape(x.shape[0], -1) - return ((flat * flat).sum(axis=-1) + 1e-20).sqrt().mean().numpy().item() - - -def normalize_per_sample(x: T) -> T: - flat = x.reshape(x.shape[0], -1) - mag = ((flat * flat).sum(axis=-1) + 1e-20).sqrt() - view_shape = (x.shape[0],) + (1,) * (len(x.shape) - 1) - return x / mag.reshape(view_shape) - - def profile( dataset: str, model_path: str, num_steps: int, batch_size: int, eps_x: float, eps_t: float ): diff --git a/experiments/homotopy_runner.py b/experiments/homotopy_runner.py new file mode 100644 index 0000000..70e1729 --- /dev/null +++ b/experiments/homotopy_runner.py @@ -0,0 +1,47 @@ +""" +Runner for the Flow Matching x Homotopy Continuation experiment plan +(see HOMOTOPY_EXPERIMENTS.md). Dispatches by experiment number to the +per-experiment script; unrecognized args are forwarded unchanged. + +All experiments log to the same MLflow experiment ("homotopy_experiments", +sqlite:///mlflow.db) so results are comparable across runs. + +Usage: + uv run experiments/homotopy_runner.py 1 + uv run experiments/homotopy_runner.py 1 --nfe-budgets 10 20 --beta 5.0 + uv run experiments/homotopy_runner.py 3 # not implemented yet +""" + +import argparse +import subprocess +import sys +from pathlib import Path + +EXPERIMENTS = { + 1: ("adaptive_sampler.py", "Condition-Number-Adaptive Sampler"), + 2: (None, "Predictor-Corrector Sampling"), + 3: (None, "Mode-Matched Base Distribution"), + 4: (None, "Implicit Homotopy Generation (IHG)"), +} + + +def main(): + parser = argparse.ArgumentParser( + description="Dispatch one experiment from HOMOTOPY_EXPERIMENTS.md" + ) + parser.add_argument("experiment", type=int, choices=sorted(EXPERIMENTS)) + args, rest = parser.parse_known_args() + + script, name = EXPERIMENTS[args.experiment] + if script is None: + print(f"Experiment {args.experiment} ({name}) is not implemented yet.") + print("See HOMOTOPY_EXPERIMENTS.md for the plan.") + sys.exit(1) + + script_path = Path(__file__).parent / script + print(f"Running experiment {args.experiment}: {name}") + subprocess.run([sys.executable, str(script_path), *rest], check=True) + + +if __name__ == "__main__": + main()