From 666e93ca084f9d42e3c467a4f5744cabb1ad1754 Mon Sep 17 00:00:00 2001 From: Ilia Ilmer Date: Sun, 26 Jul 2026 05:56:31 -0400 Subject: [PATCH 1/4] reorg files --- .github/workflows/lint.yml | 4 +- .pre-commit-config.yaml | 17 +- README.md | 40 +-- examples/generate_animation.py | 295 ----------------- examples/generate_moons.py | 253 --------------- examples/generate_samples.py | 142 -------- examples/train_images_hydra.py | 287 ---------------- examples/train_moons_hydra.py | 248 -------------- experiments/_common.py | 45 +++ experiments/curvature_profile.py | 294 +++++++++++++++++ experiments/schedule_compare.py | 138 ++++++++ experiments/solver_compare.py | 183 +++++++++++ scripts/_common.py | 128 ++++++++ scripts/download_datasets.py | 2 - .../generate_images.py => scripts/generate.py | 305 +++++++++++++----- scripts/train.py | 268 +++++++++++++++ 16 files changed, 1306 insertions(+), 1343 deletions(-) delete mode 100644 examples/generate_animation.py delete mode 100644 examples/generate_moons.py delete mode 100644 examples/generate_samples.py delete mode 100644 examples/train_images_hydra.py delete mode 100644 examples/train_moons_hydra.py create mode 100644 experiments/_common.py create mode 100644 experiments/curvature_profile.py create mode 100644 experiments/schedule_compare.py create mode 100644 experiments/solver_compare.py create mode 100644 scripts/_common.py rename examples/generate_images.py => scripts/generate.py (63%) create mode 100644 scripts/train.py diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 4e633de..1ed655b 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -31,12 +31,12 @@ jobs: - name: Run ruff check (linting) run: | - uv run ruff check tinyflow/ tests/ examples/ --output-format=github + uv run ruff check tinyflow/ tests/ scripts/ experiments/ --output-format=github continue-on-error: true - name: Run ruff format (formatting check) run: | - uv run ruff format --check tinyflow/ tests/ examples/ + uv run ruff format --check tinyflow/ tests/ scripts/ experiments/ continue-on-error: true - name: Run pyright (type checking) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 88c2b27..e27c823 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,7 +3,7 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.8.4 + rev: v0.16.0 hooks: # Run the linter - id: ruff @@ -29,18 +29,3 @@ repos: - id: end-of-file-fixer # Remove trailing whitespace - id: trailing-whitespace - # Prevent committing to main branch directly (optional - uncomment if desired) - # - id: no-commit-to-branch - # args: ['--branch', 'main'] - - # Optional: Type checking with pyright (can be slow) - # Uncomment if you want type checking on every commit - # - repo: https://github.com/RobertCraigie/pyright-python - # rev: v1.1.390 - # hooks: - # - id: pyright - # additional_dependencies: [ - # 'tinygrad>=0.10.0', - # 'numpy>=2.2.1', - # 'loguru>=0.7.3', - # ] diff --git a/README.md b/README.md index 6e662b3..f798c4a 100644 --- a/README.md +++ b/README.md @@ -52,37 +52,37 @@ uv sync ```bash # Train on 2D moons dataset (fast, for testing) -uv run examples/train_moons_hydra.py +uv run scripts/train.py dataset=moons model=neural_network training=moons_default scheduler=linear # Train on MNIST -uv run examples/train_mnist_hydra.py +uv run scripts/train.py # Train on Fashion MNIST -uv run examples/train_mnist_hydra.py dataset=fashion_mnist +uv run scripts/train.py dataset=fashion_mnist # Train on CIFAR-10 -uv run examples/train_mnist_hydra.py dataset=cifar10 +uv run scripts/train.py dataset=cifar10 model=unet ``` ### Generating Samples ```bash # Generate static grid for MNIST -uv run examples/generate_images.py generation.model_path=model_mnist_unet_linear.safetensors +uv run scripts/generate.py generation.model_path=model_mnist_unet_linear.safetensors # Generate with class predictions and FID score -uv run examples/generate_images.py \ +uv run scripts/generate.py \ generation.model_path=model_mnist_unet_linear.safetensors \ generation.show_predictions=true \ generation.compute_fid=true # Generate animated GIF showing the flow process -uv run examples/generate_images.py \ +uv run scripts/generate.py \ generation.model_path=model_mnist_unet_linear.safetensors \ --animated -# Generate for 2D moons dataset -uv run examples/generate_moons.py generation.model_path=model_moons_neural_network_linear.safetensors +# Generate for 2D moons dataset (switch the base config) +uv run scripts/generate.py --config-name=generate_moons_config generation.model_path=model_moons_neural_network_linear.safetensors ``` ### Training FID Classifier (for evaluation) @@ -133,14 +133,16 @@ configs/ # Hydra configuration files ├── training/ # Training parameters └── generation/ # Generation settings -examples/ -├── train_mnist_hydra.py # MNIST/CIFAR-10 training -├── train_moons_hydra.py # 2D moons training -├── generate_images.py # Image generation with predictions/FID -└── generate_moons.py # 2D moons generation - scripts/ -└── train_fid_classifier.py # Train FID classifier with validation +├── train.py # Unified training entry point (moons + image datasets) +├── generate.py # Unified generation entry point (predictions/FID/animation) +├── train_fid_classifier.py # Train FID classifier with validation +└── download_datasets.py # Dataset download helper + +experiments/ # ODE-solver and time-schedule research scripts +├── solver_compare.py # Compare solvers (Euler/Heun/RK4) at fixed NFE +├── schedule_compare.py # Compare uniform vs back-loaded time schedules +└── curvature_profile.py # Velocity-field curvature/LTE-proxy profiling docs/ ├── MLFLOW_HYDRA_GUIDE.md # Experiment tracking guide @@ -155,13 +157,13 @@ This project uses Hydra for flexible configuration management: ```bash # Override specific parameters -uv run examples/train_mnist_hydra.py scheduler=cosine optimizer.lr=0.001 +uv run scripts/train.py scheduler=cosine optimizer.lr=0.001 # Compare multiple schedulers (multirun) -uv run examples/train_mnist_hydra.py -m scheduler=linear,cosine,polynomial +uv run scripts/train.py -m scheduler=linear,cosine,polynomial # Use experiment configs -uv run examples/train_moons_hydra.py +experiment=quick_test +uv run scripts/train.py +experiment=quick_test ``` See `docs/MLFLOW_HYDRA_GUIDE.md` for detailed usage. diff --git a/examples/generate_animation.py b/examples/generate_animation.py deleted file mode 100644 index 20f0dad..0000000 --- a/examples/generate_animation.py +++ /dev/null @@ -1,295 +0,0 @@ -""" -Generate an animated GIF showing the flow matching generation process. - -Usage: - uv run examples/generate_animation.py --model model.safetensors --dataset mnist --grid-size 3 - uv run examples/generate_animation.py --model model.safetensors --dataset fashion_mnist --grid-size 4 --fps 10 - uv run examples/generate_animation.py --model model.safetensors --dataset mnist --per-frame-norm # Show noise structure - uv run examples/generate_animation.py --model model.safetensors --dataset mnist --show-distribution # Show distribution evolution -""" - -import argparse - -import matplotlib.pyplot as plt -import numpy as np -from PIL import Image -from tinygrad import TinyJit -from tinygrad.tensor import Tensor as T -from tqdm import tqdm - -from tinyflow.nn import UNetCIFAR10Large, UNetTinygrad -from tinyflow.solver import RK4 -from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist - - -def generate_animation( - model_path: str, - dataset: str = "mnist", - model_type: str = "unet", - grid_size: int = 3, - num_steps: int = 100, - num_frames: int = 50, - fps: int = 10, - seed: int = 42, - per_frame_norm: bool = False, - show_distribution: bool = False, -): - """Generate an animated GIF of the generation process.""" - - # Set seed - T.manual_seed(seed) - - # Load model - print(f"Loading model from {model_path}...") - from tinyflow.trainer import BaseTrainer - - if model_type == "unet": - from tinygrad.nn.state import safe_load as _safe_load - - _state = _safe_load(model_path) - _in_ch = 1 if dataset in ["mnist", "fashion_mnist"] else 3 - _time_embed_dim = ( - int(_state["enc1.conv.weight"].shape[1]) - _in_ch - if "enc1.conv.weight" in _state - else 64 - ) - model = UNetTinygrad(time_embed_dim=_time_embed_dim) - elif model_type == "unet_large": - if dataset != "cifar10": - raise ValueError("unet_large only supports cifar10") - model = UNetCIFAR10Large() - else: - raise ValueError(f"Unknown model type: {model_type}") - - BaseTrainer.load_model(model, model_path) - - # Configure for dataset - if dataset in ["mnist", "fashion_mnist"]: - shape = (grid_size * grid_size, 1, 28, 28) - preprocess_hook = preprocess_time_mnist - is_color = False - elif dataset == "cifar10": - shape = (grid_size * grid_size, 3, 32, 32) - preprocess_hook = preprocess_time_cifar - is_color = True - else: - raise ValueError(f"Unknown dataset: {dataset}") - - # Create solver - solver = RK4(model, preprocess_hook=preprocess_hook) - - # Generate samples and capture frames - print(f"Generating {grid_size}×{grid_size} samples with {num_frames} frames...") - T.training = False - - x = T.randn(*shape) - h_step = 1.0 / num_steps - - # Determine which steps to capture - capture_steps = np.linspace(0, num_steps - 1, num_frames, dtype=int) - snapshots = [] # (step_index, numpy_array) pairs - - # Capture the initial noise as the first frame - snapshots.append((0, x.numpy().copy())) - - # JIT compile the solver step for better performance - @TinyJit - def jit_step(h, t, x): - return solver.sample(h, t, x) - - # Solve ODE and capture intermediate states - for step in tqdm(range(num_steps), desc="Generating"): - t = (T.zeros(1) + step * h_step).contiguous() - x = jit_step(h_step, t, x) - - if step in capture_steps: - # Ensure tensor is realized before capturing to avoid memory accumulation - x.realize() - snapshots.append((step + 1, x.numpy().copy())) - - # Compute normalization parameters - if per_frame_norm: - print("Using per-frame normalization (shows noise structure)") - global_min = global_max = None - else: - print("Using global normalization (consistent brightness)") - global_min = min(s.min() for _, s in snapshots) - global_max = max(s.max() for _, s in snapshots) - - # Store raw snapshots for distribution visualization - if show_distribution: - initial_raw = snapshots[0][1] # Raw initial noise - final_raw = snapshots[-1][1] # Raw final images - - # Render frames - frames = [] - for step_idx, x_np in snapshots: - if per_frame_norm: - # Normalize each frame independently - frame_min = x_np.min() - frame_max = x_np.max() - x_normalized = (x_np - frame_min) / (frame_max - frame_min + 1e-8) - else: - # Normalize all frames to same global scale - x_normalized = (x_np - global_min) / (global_max - global_min + 1e-8) - - x_normalized = np.clip(x_normalized, 0, 1) - - # Create figure with optional distribution panel - if show_distribution: - fig = plt.figure(figsize=(14, 8)) - gs = fig.add_gridspec( - grid_size, - grid_size + 1, - width_ratios=[1] * grid_size + [1.2], - wspace=0.3, - hspace=0.1, - ) - axes = [[fig.add_subplot(gs[i, j]) for j in range(grid_size)] for i in range(grid_size)] - ax_dist = fig.add_subplot(gs[:, -1]) - else: - fig, axes = plt.subplots(grid_size, grid_size, figsize=(8, 8)) - if grid_size == 1: - axes = [[axes]] - - fig.suptitle(f"t = {step_idx / num_steps:.2f}", fontsize=16, y=0.98) - - # Render image grid - for i in range(grid_size): - for j in range(grid_size): - idx = i * grid_size + j - ax = axes[i][j] - - if is_color: - img = x_normalized[idx].transpose(1, 2, 0) - ax.imshow(img) - else: - img = x_normalized[idx, 0] - ax.imshow(img, cmap="gray") - - ax.axis("off") - - # Add distribution visualization - if show_distribution: - ax_dist.clear() - - # Plot histograms of raw pixel values (before normalization) - initial_flat = initial_raw.flatten() - final_flat = final_raw.flatten() - current_flat = x_np.flatten() - - # Compute common bins for consistent visualization - all_values = np.concatenate([initial_flat, final_flat, current_flat]) - bins = np.linspace(all_values.min(), all_values.max(), 50) - - # Plot reference distributions with transparency - ax_dist.hist( - initial_flat, - bins=bins, - alpha=0.3, - color="blue", - label="t=0.00 (noise)", - density=True, - ) - ax_dist.hist( - final_flat, bins=bins, alpha=0.3, color="green", label="t=1.00 (data)", density=True - ) - - # Plot current distribution prominently - ax_dist.hist( - current_flat, - bins=bins, - alpha=0.7, - color="red", - label=f"t={step_idx / num_steps:.2f}", - density=True, - ) - - ax_dist.set_xlabel("Pixel Value") - ax_dist.set_ylabel("Density") - ax_dist.set_title("Distribution Evolution") - ax_dist.legend(fontsize=8) - ax_dist.grid(True, alpha=0.3) - - plt.tight_layout() - - fig.canvas.draw() - data_rgba = np.asarray(fig.canvas.buffer_rgba()) - frame = data_rgba[..., :3] - frames.append(Image.fromarray(frame)) - plt.close(fig) - - # Save as GIF - output_file = f"generation_{dataset}_{grid_size}x{grid_size}.gif" - print(f"Saving animation to {output_file}...") - - # Add final frame multiple times to pause at the end - frames.extend([frames[-1]] * fps) - - frames[0].save( - output_file, - save_all=True, - append_images=frames[1:], - duration=1000 // fps, # milliseconds per frame - loop=0, - ) - - print(f"✓ Animation saved to {output_file}") - print(f" Frames: {len(frames)}") - print(f" FPS: {fps}") - print(f" Duration: {len(frames) / fps:.1f} seconds") - - -def main(): - parser = argparse.ArgumentParser(description="Generate animation of flow matching generation") - parser.add_argument( - "--model", type=str, default="model.safetensors", help="Path to model weights" - ) - parser.add_argument( - "--dataset", - type=str, - default="mnist", - choices=["mnist", "fashion_mnist", "cifar10"], - help="Dataset type", - ) - parser.add_argument( - "--model-type", - type=str, - default="unet", - choices=["unet", "unet_large"], - help="Model architecture type", - ) - parser.add_argument("--grid-size", type=int, default=3, help="Grid size (n×n)") - parser.add_argument("--num-steps", type=int, default=100, help="Number of ODE steps") - parser.add_argument("--num-frames", type=int, default=50, help="Number of frames in animation") - parser.add_argument("--fps", type=int, default=10, help="Frames per second") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument( - "--per-frame-norm", - action="store_true", - help="Use per-frame normalization to show noise structure (vs global normalization)", - ) - parser.add_argument( - "--show-distribution", - action="store_true", - help="Add histogram showing evolution of pixel value distribution", - ) - - args = parser.parse_args() - - generate_animation( - model_path=args.model, - dataset=args.dataset, - model_type=args.model_type, - grid_size=args.grid_size, - num_steps=args.num_steps, - num_frames=args.num_frames, - fps=args.fps, - seed=args.seed, - per_frame_norm=args.per_frame_norm, - show_distribution=args.show_distribution, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/generate_moons.py b/examples/generate_moons.py deleted file mode 100644 index 7a917df..0000000 --- a/examples/generate_moons.py +++ /dev/null @@ -1,253 +0,0 @@ -""" -Generate samples from a trained flow matching model for the 2D moons dataset. - -Usage: - # Generate static visualization - uv run examples/generate_moons.py generation.model_path=model_moons_neural_network_linear.safetensors - - # Generate with animation - uv run examples/generate_moons.py generation.model_path=model_moons_neural_network_linear.safetensors --animated - - # Override number of samples and steps - uv run examples/generate_moons.py generation.model_path=model.safetensors generation.n_samples=200 generation.num_steps=200 -""" - -import os - -import hydra -import matplotlib.pyplot as plt -import numpy as np -from omegaconf import DictConfig, OmegaConf -from PIL import Image -from tinygrad import TinyJit -from tinygrad.tensor import Tensor as T -from tqdm import tqdm - -from tinyflow.nn import NeuralNetwork -from tinyflow.solver import DDIM, Euler, Heun, MidpointSolver, RK4 -from tinyflow.trainer import BaseTrainer -from tinyflow.utils import preprocess_time_moons - -plt.style.use("ggplot") - - -def create_solver(cfg: DictConfig, model, preprocess_hook): - """Create ODE solver from config.""" - solver_type = cfg.solver.type - if solver_type == "euler": - return Euler(model, preprocess_hook=preprocess_hook) - if solver_type == "heun": - return Heun(model, preprocess_hook=preprocess_hook) - if solver_type == "midpoint": - return MidpointSolver(model, preprocess_hook=preprocess_hook) - if solver_type == "rk4": - return RK4(model, preprocess_hook=preprocess_hook) - if solver_type == "ddim": - eta = cfg.solver.get("eta", 0.0) - return DDIM(model, preprocess_hook=preprocess_hook, eta=eta) - raise ValueError(f"Unknown solver type: {solver_type}") - - -def create_model(cfg: DictConfig): - """Create model from config.""" - model_type = cfg.model.type - if model_type == "neural_network": - return NeuralNetwork( - in_dim=cfg.model.input_dim, - time_embed_dim=cfg.model.time_embed_dim, - out_dim=cfg.model.output_dim, - ) - raise ValueError(f"Unknown model type: {model_type}") - - -def generate_static_visualization(cfg: DictConfig, solver): - """Generate static visualization with multiple time snapshots.""" - n_samples = cfg.generation.n_samples - num_steps = cfg.generation.num_steps - step_size = cfg.generation.step_size - num_plots = cfg.generation.num_plots - - print(f"Generating {n_samples} samples with {num_plots} snapshots...") - - # Initialize samples - T.training = False - x = T.randn(n_samples, 2) - time_grid = T.linspace(0, 1, num_steps) - sample_every = time_grid.shape[0] // num_plots - - # Store snapshots to visualize - snapshots = [] - snapshot_times = [] - - # JIT compile the solver step for better performance - @TinyJit - def jit_step(step_size, t, x): - return solver.sample(step_size, t, x) - - # Generate all samples first - for idx in tqdm(range(int(time_grid.shape[0])), desc="Generating samples"): - t = time_grid[idx].contiguous() - x = jit_step(step_size, t, x) - - # Store reference for visualization - if (idx + 1) % sample_every == 0: - snapshots.append(x) - snapshot_times.append(t) - - # Visualize all snapshots - _, ax = plt.subplots(1, num_plots, figsize=(30, 4), sharex=True, sharey=True) - for i, (snapshot, t) in enumerate(zip(snapshots, snapshot_times, strict=False)): - x_np = snapshot.numpy() - ax[i].scatter(x_np[:, 0], x_np[:, 1], s=20, alpha=0.6) - ax[i].set_title(f"t={t.numpy():.2f}") - ax[i].set_xlim(-3, 3) - ax[i].set_ylim(-3, 3) - ax[i].grid(True, alpha=0.3) - - plt.tight_layout() - - # Save - os.makedirs(cfg.generation.output_dir, exist_ok=True) - output_file = os.path.join(cfg.generation.output_dir, "moons_generation.png") - plt.savefig(output_file, dpi=150, bbox_inches="tight") - print(f"Saved to {output_file}") - - if cfg.generation.get("show_plot", False): - plt.show() - else: - plt.close() - - -def generate_animation(cfg: DictConfig, solver): - """Generate an animated GIF of the generation process.""" - n_samples = cfg.generation.n_samples - num_steps = cfg.generation.num_steps - step_size = cfg.generation.step_size - num_plots = cfg.generation.num_plots - fps = cfg.generation.fps - - print(f"Generating {n_samples} samples with {num_plots} frames...") - - # Initialize samples - T.training = False - x = T.randn(n_samples, 2) - time_grid = T.linspace(0, 1, num_steps) - sample_every = time_grid.shape[0] // num_plots - - # Store snapshots to visualize - snapshots = [] - snapshot_times = [] - - # JIT compile the solver step for better performance - @TinyJit - def jit_step(step_size, t, x): - return solver.sample(step_size, t, x) - - # Generate all samples first - for idx in tqdm(range(int(time_grid.shape[0])), desc="Generating animation"): - t = time_grid[idx].contiguous() - x = jit_step(step_size, t, x) - - # Store reference for visualization - if (idx + 1) % sample_every == 0: - snapshots.append(x) - snapshot_times.append(t) - - # Create frames - frames = [] - for snapshot, t in tqdm( - zip(snapshots, snapshot_times, strict=False), desc="Creating frames", total=len(snapshots) - ): - x_np = snapshot.numpy() - - fig, ax = plt.subplots(figsize=(8, 8)) - ax.scatter(x_np[:, 0], x_np[:, 1], s=20, alpha=0.6) - ax.set_title(f"t = {t.numpy():.2f}", fontsize=16) - ax.set_xlim(-3, 3) - ax.set_ylim(-3, 3) - ax.grid(True, alpha=0.3) - ax.set_aspect("equal") - - # Convert plot to image - fig.canvas.draw() - # Use buffer_rgba() for compatibility across backends - buf = fig.canvas.buffer_rgba() - frame = np.asarray(buf) - # Convert RGBA to RGB - frame = frame[:, :, :3] - frames.append(Image.fromarray(frame)) - plt.close(fig) - - # Save as GIF - os.makedirs(cfg.generation.output_dir, exist_ok=True) - output_file = os.path.join(cfg.generation.output_dir, "moons_generation.gif") - - # Add final frame multiple times to pause at the end - frames.extend([frames[-1]] * fps) - - frames[0].save( - output_file, - save_all=True, - append_images=frames[1:], - duration=1000 // fps, - loop=0, - ) - - print(f"Animation saved to {output_file}") - print(f" Frames: {len(frames)}") - print(f" FPS: {fps}") - print(f" Duration: {len(frames) / fps:.1f} seconds") - - -def main_impl(cfg: DictConfig): - """Main generation function.""" - animated = cfg.generation.get("animated", False) - - print("Configuration:") - print(OmegaConf.to_yaml(cfg)) - print(f"\nAnimated: {animated}") - - # Set random seed - if cfg.get("seed"): - T.manual_seed(cfg.seed) - - # Load model - model_path = cfg.generation.model_path - if not os.path.exists(model_path): - raise FileNotFoundError(f"Model file not found: {model_path}") - - print(f"\nLoading model from {model_path}...") - model = create_model(cfg) - BaseTrainer.load_model(model, model_path) - - # Create solver - # Note: Scheduler/path not needed at generation time - the model has already - # learned the velocity field during training - solver = create_solver(cfg, model, preprocess_time_moons) - - # Generate samples - if animated: - generate_animation(cfg, solver) - else: - generate_static_visualization(cfg, solver) - - -@hydra.main(version_base=None, config_path="../configs", config_name="generate_moons_config") -def main(cfg: DictConfig): - """Hydra entry point.""" - main_impl(cfg) - - -if __name__ == "__main__": - import sys - - # Check for --animated flag and remove it before Hydra processes args - animated = "--animated" in sys.argv - if animated: - sys.argv.remove("--animated") - - # Pass animated flag through config override - if animated: - sys.argv.append("generation.animated=true") - - main() diff --git a/examples/generate_samples.py b/examples/generate_samples.py deleted file mode 100644 index ec1b2be..0000000 --- a/examples/generate_samples.py +++ /dev/null @@ -1,142 +0,0 @@ -""" -Generate samples from a trained flow matching model. - -Usage: - uv run examples/generate_samples.py --model model.safetensors --dataset mnist --grid-size 5 - uv run examples/generate_samples.py --model model.safetensors --dataset fashion_mnist --grid-size 4 - uv run examples/generate_samples.py --model model.safetensors --dataset cifar10 --grid-size 3 -""" - -import argparse - -import matplotlib.pyplot as plt -import numpy as np -from tinygrad import TinyJit -from tinygrad.tensor import Tensor as T - -from tinyflow.nn import UNetTinygrad -from tinyflow.solver import RK4 -from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist - - -def generate_grid( - model_path: str, - dataset: str = "mnist", - grid_size: int = 5, - num_steps: int = 100, - seed: int = 42, -): - """Generate a grid of samples from a trained model.""" - - # Set seed - T.manual_seed(seed) - - # Load model - print(f"Loading model from {model_path}...") - model = UNetTinygrad() - from tinyflow.trainer import BaseTrainer - - BaseTrainer.load_model(model, model_path) - - # Configure for dataset - if dataset in ["mnist", "fashion_mnist"]: - shape = (grid_size * grid_size, 1, 28, 28) - preprocess_hook = preprocess_time_mnist - is_color = False - elif dataset == "cifar10": - shape = (grid_size * grid_size, 3, 32, 32) - preprocess_hook = preprocess_time_cifar - is_color = True - else: - raise ValueError(f"Unknown dataset: {dataset}") - - # Create solver - solver = RK4(model, preprocess_hook=preprocess_hook) - - # Generate samples - print(f"Generating {grid_size}×{grid_size} = {grid_size * grid_size} samples...") - T.training = False - - x = T.randn(*shape) - h_step = 1.0 / num_steps - - # JIT compile the solver step for better performance - @TinyJit - def jit_step(h, t, x): - return solver.sample(h, t, x) - - # Solve ODE from t=0 to t=1 - for step in range(num_steps): - t = (T.zeros(1) + step * h_step).contiguous() - x = jit_step(h_step, t, x) - - # Ensure final result is realized before conversion - x.realize() - - # Convert to numpy and normalize - x_np = x.numpy() - - # Normalize for display - x_np = (x_np - x_np.min()) / (x_np.max() - x_np.min() + 1e-8) - x_np = np.clip(x_np, 0, 1) - - # Plot grid - fig, axes = plt.subplots(grid_size, grid_size, figsize=(10, 10)) - - for i in range(grid_size): - for j in range(grid_size): - idx = i * grid_size + j - ax = axes[i, j] if grid_size > 1 else axes - - if is_color: - # CIFAR-10: (C, H, W) -> (H, W, C) - img = x_np[idx].transpose(1, 2, 0) - ax.imshow(img) - else: - # MNIST/Fashion MNIST: (1, H, W) -> (H, W) - img = x_np[idx, 0] - ax.imshow(img, cmap="gray") - - ax.axis("off") - - plt.tight_layout() - - # Save - output_file = f"generated_{dataset}_{grid_size}x{grid_size}.png" - plt.savefig(output_file, dpi=150, bbox_inches="tight") - print(f"✓ Saved to {output_file}") - - plt.show() - - -def main(): - parser = argparse.ArgumentParser( - description="Generate samples from trained flow matching model" - ) - parser.add_argument( - "--model", type=str, default="model.safetensors", help="Path to model weights" - ) - parser.add_argument( - "--dataset", - type=str, - default="mnist", - choices=["mnist", "fashion_mnist", "cifar10"], - help="Dataset type", - ) - parser.add_argument("--grid-size", type=int, default=5, help="Grid size (n×n)") - parser.add_argument("--num-steps", type=int, default=100, help="Number of ODE steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - - args = parser.parse_args() - - generate_grid( - model_path=args.model, - dataset=args.dataset, - grid_size=args.grid_size, - num_steps=args.num_steps, - seed=args.seed, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/train_images_hydra.py b/examples/train_images_hydra.py deleted file mode 100644 index c64cf8d..0000000 --- a/examples/train_images_hydra.py +++ /dev/null @@ -1,287 +0,0 @@ -""" -Image dataset training script with Hydra configuration and MLflow tracking. -Supports MNIST, Fashion MNIST, and CIFAR-10 datasets. -""" - -import os - -import hydra -import matplotlib.pyplot as plt -import mlflow -from omegaconf import DictConfig, OmegaConf -from tinygrad.nn.optim import Adam -from tinygrad.nn.state import get_parameters -from tinygrad.tensor import Tensor as T - -from tinyflow.dataloader import CIFAR10Loader, FashionMNISTLoader, MNISTLoader -from tinyflow.losses import mse -from tinyflow.nn import UNetCIFAR10, UNetCIFAR10Large, UNetMNIST -from tinyflow.nn_utils.lr_scheduler import ( - CosineAnnealingLR, - StepLRScheduler, - WarmupScheduler, -) -from tinyflow.path import AffinePath -from tinyflow.path.scheduler import ( - CosineScheduler, - LinearScheduler, - LinearVarPresScheduler, - PolynomialScheduler, -) -from tinyflow.solver import DDIM, Euler, Heun, MidpointSolver, RK4 -from tinyflow.trainer import CIFAR10Trainer, MNISTTrainer -from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist - -plt.style.use("ggplot") - -mlflow.set_tracking_uri("sqlite:///mlflow.db") -mlflow.set_experiment("flow_matching") - - -def create_scheduler(cfg: DictConfig): - """Create scheduler from config.""" - scheduler_type = cfg.scheduler.type - if scheduler_type == "LinearScheduler": - return LinearScheduler() - if scheduler_type == "CosineScheduler": - return CosineScheduler() - if scheduler_type == "LinearVarPresScheduler": - return LinearVarPresScheduler() - if scheduler_type == "PolynomialScheduler": - degree = cfg.scheduler.get("degree", 2) - return PolynomialScheduler(degree) - raise ValueError(f"Unknown scheduler type: {scheduler_type}") - - -def create_lr_scheduler(cfg: DictConfig, optimizer): - """Create learning rate scheduler from config.""" - if not cfg.get("lr_scheduler"): - return None - - scheduler_type = cfg.lr_scheduler.type - - # Create base scheduler - if scheduler_type == "NullLRScheduler": - return None - elif scheduler_type == "StepLRScheduler": - base_scheduler = StepLRScheduler( - optimizer, - step_size=cfg.lr_scheduler.get("step_size", 1000), - gamma=cfg.lr_scheduler.get("gamma", 0.1), - ) - elif scheduler_type == "CosineAnnealingLR": - base_scheduler = CosineAnnealingLR( - optimizer, - t_max=cfg.lr_scheduler.get("t_max", 5000), - eta_min=cfg.lr_scheduler.get("eta_min", 0.0), - warm=cfg.lr_scheduler.get("warm", False), - ) - elif scheduler_type == "WarmupScheduler": - # Create nested base scheduler - base_cfg = cfg.lr_scheduler.base_scheduler - if base_cfg.type == "StepLRScheduler": - nested_scheduler = StepLRScheduler( - optimizer, - step_size=base_cfg.get("step_size", 1000), - gamma=base_cfg.get("gamma", 0.1), - ) - elif base_cfg.type == "CosineAnnealingLR": - nested_scheduler = CosineAnnealingLR( - optimizer, - t_max=base_cfg.get("t_max", 5000), - eta_min=base_cfg.get("eta_min", 0.0), - warm=base_cfg.get("warm", False), - ) - else: - raise ValueError(f"Unknown base scheduler type: {base_cfg.type}") - - return WarmupScheduler( - optimizer, - base_scheduler=nested_scheduler, - warmup_steps=cfg.lr_scheduler.get("warmup_steps", 500), - warmup_start_lr=cfg.lr_scheduler.get("warmup_start_lr", 0.0), - ) - else: - raise ValueError(f"Unknown LR scheduler type: {scheduler_type}") - - return base_scheduler - - -def create_model(cfg: DictConfig): - """Create model from config.""" - model_type = cfg.model.type - dataset_type = cfg.dataset.get("type", cfg.dataset.name) - - if model_type == "unet": - if dataset_type in ["mnist", "fashion_mnist"]: - return UNetMNIST(in_channels=1, out_channels=1) - elif dataset_type == "cifar10": - return UNetCIFAR10(in_channels=3, out_channels=3) - elif model_type == "unet_large": - if dataset_type == "cifar10": - return UNetCIFAR10Large(in_channels=3, out_channels=3) - else: - raise ValueError(f"unet_large only supports cifar10, got {dataset_type}") - - raise ValueError(f"Unknown model type: {model_type} with dataset type: {dataset_type}") - - -def create_dataloader(cfg: DictConfig): - """Create dataloader from config.""" - dataset_type = cfg.dataset.get("type", cfg.dataset.name) - batch_size = cfg.dataset.get("batch_size", 32) - flatten = cfg.dataset.get("flatten", False) - shuffle = cfg.dataset.get("shuffle", True) - - if dataset_type == "mnist": - return MNISTLoader(flatten=flatten, batch_size=batch_size, shuffle=shuffle) - elif dataset_type == "fashion_mnist": - path = cfg.dataset.get("path", "dataset/fashion_mnist") - train = cfg.dataset.get("train", True) - return FashionMNISTLoader( - path=path, flatten=flatten, batch_size=batch_size, shuffle=shuffle, train=train - ) - elif dataset_type == "cifar10": - path = cfg.dataset.get("path", "dataset/cifar10/cifar-10-batches-py") - cache = cfg.dataset.get("cache", True) - normalize = cfg.dataset.get("normalize", True) - train = cfg.dataset.get("train", True) - return CIFAR10Loader( - path=path, - batch_size=batch_size, - shuffle=shuffle, - train=train, - cache=cache, - normalize=normalize, - ) - else: - raise ValueError(f"Unknown dataset type: {dataset_type}") - - -def create_trainer(cfg: DictConfig, model, dataloader, optim, path, lr_scheduler=None): - """Create trainer from config.""" - dataset_type = cfg.dataset.get("type", cfg.dataset.name) - num_epochs = cfg.training.num_epochs - log_interval = cfg.training.log_interval - gradient_accumulation_steps = cfg.training.get("gradient_accumulation_steps", 1) - - if dataset_type in ["mnist", "fashion_mnist"]: - return MNISTTrainer( - model=model, - dataloader=dataloader, - optim=optim, - loss_fn=mse, - path=path, - num_epochs=num_epochs, - log_interval=log_interval, - lr_scheduler=lr_scheduler, - gradient_accumulation_steps=gradient_accumulation_steps, - ) - elif dataset_type == "cifar10": - return CIFAR10Trainer( - model=model, - dataloader=dataloader, - optim=optim, - loss_fn=mse, - path=path, - num_epochs=num_epochs, - log_interval=log_interval, - lr_scheduler=lr_scheduler, - gradient_accumulation_steps=gradient_accumulation_steps, - ) - else: - raise ValueError(f"Unknown dataset type: {dataset_type}") - - -def create_solver(cfg: DictConfig, model, preprocess_hook): - """Create ODE solver from config.""" - solver_type = cfg.solver.type - if solver_type == "euler": - return Euler(model, preprocess_hook=preprocess_hook) - if solver_type == "heun": - return Heun(model, preprocess_hook=preprocess_hook) - if solver_type == "midpoint": - return MidpointSolver(model, preprocess_hook=preprocess_hook) - if solver_type == "rk4": - return RK4(model, preprocess_hook=preprocess_hook) - if solver_type == "ddim": - eta = cfg.solver.get("eta", 0.0) - return DDIM(model, preprocess_hook=preprocess_hook, eta=eta) - raise ValueError(f"Unknown solver type: {solver_type}") - - -def get_preprocess_hook(cfg: DictConfig): - """Get preprocessing hook for time input based on dataset.""" - dataset_type = cfg.dataset.get("type", cfg.dataset.name) - if dataset_type in ["mnist", "fashion_mnist"]: - return preprocess_time_mnist - elif dataset_type == "cifar10": - return preprocess_time_cifar - else: - raise ValueError(f"Unknown dataset type: {dataset_type}") - - -def generate_model_name(cfg: DictConfig) -> str: - """Generate model name from config.""" - # Get components - dataset = cfg.dataset.get("type", cfg.dataset.name) - model_type = cfg.model.type - scheduler = cfg.scheduler.type.replace("Scheduler", "").lower() - - # Build name: model___.safetensors - name = f"model_{dataset}_{model_type}_{scheduler}.safetensors" - return name - - -@hydra.main(version_base=None, config_path="../configs", config_name="config") -def main(cfg: DictConfig): - """Main training function.""" - print("Configuration:") - print(OmegaConf.to_yaml(cfg)) - - # Set random seed if specified - if cfg.get("seed"): - T.manual_seed(cfg.seed) - - # Generate model name (can be overridden with model_name parameter) - model_name = cfg.get("model_name", generate_model_name(cfg)) - print(f"\nModel will be saved as: {model_name}") - - # Create model, scheduler, and path - model = create_model(cfg) - scheduler = create_scheduler(cfg) - path = AffinePath(scheduler=scheduler) - - # Create optimizer - optim = Adam(get_parameters(model), lr=cfg.optimizer.lr) - - # Create learning rate scheduler - lr_scheduler = create_lr_scheduler(cfg, optim) - - # Create dataloader and trainer based on dataset type - dataloader = create_dataloader(cfg) - trainer = create_trainer(cfg, model, dataloader, optim, path, lr_scheduler) - - # Train the model - dataset_name = cfg.dataset.get("type", cfg.dataset.name) - with mlflow.start_run(run_name=dataset_name): - mlflow.log_params(dict(cfg)) - mlflow.log_param("model_name", model_name) - model = trainer.train() - trainer.save_model(model_name) - - # Save loss plot - if cfg.training.get("log_artifacts", True): - output_dir = cfg.get("output_dir", "outputs") - os.makedirs(output_dir, exist_ok=True) - trainer.plot_loss(output_dir, log_to_mlflow=True) - - # Generate samples if requested - if cfg.training.get("generate_samples", True): - preprocess_hook = get_preprocess_hook(cfg) - solver = create_solver(cfg, model, preprocess_hook) - trainer.predict(cfg, solver) - - -if __name__ == "__main__": - main() diff --git a/examples/train_moons_hydra.py b/examples/train_moons_hydra.py deleted file mode 100644 index 5deb4b1..0000000 --- a/examples/train_moons_hydra.py +++ /dev/null @@ -1,248 +0,0 @@ -import os - -import hydra -import matplotlib.pyplot as plt -from omegaconf import DictConfig, OmegaConf -from sklearn.datasets import make_moons -from tinygrad.nn.optim import Adam -from tinygrad.nn.state import get_parameters -from tinygrad.tensor import Tensor as T -from tqdm.auto import tqdm - -from tinyflow.losses import mse -from tinyflow.nn import NeuralNetwork -from tinyflow.nn_utils.lr_scheduler import ( - CosineAnnealingLR, - StepLRScheduler, - WarmupScheduler, -) -from tinyflow.path import AffinePath -from tinyflow.path.scheduler import ( - CosineScheduler, - LinearScheduler, - LinearVarPresScheduler, - PolynomialScheduler, -) -from tinyflow.solver import DDIM, Euler, Heun, MidpointSolver, RK4 -from tinyflow.utils import preprocess_time_moons, visualize_moons - -plt.style.use("ggplot") - - -def create_scheduler(cfg: DictConfig): - """Create scheduler from config.""" - scheduler_type = cfg.scheduler.type - if scheduler_type == "LinearScheduler": - return LinearScheduler() - if scheduler_type == "CosineScheduler": - return CosineScheduler() - if scheduler_type == "LinearVarPresScheduler": - return LinearVarPresScheduler() - if scheduler_type == "PolynomialScheduler": - degree = cfg.scheduler.get("degree", 2) - return PolynomialScheduler(degree) - raise ValueError(f"Unknown scheduler type: {scheduler_type}") - - -def create_lr_scheduler(cfg: DictConfig, optimizer): - """Create learning rate scheduler from config.""" - if not cfg.get("lr_scheduler"): - return None - - scheduler_type = cfg.lr_scheduler.type - - # Create base scheduler - if scheduler_type == "NullLRScheduler": - return None - elif scheduler_type == "StepLRScheduler": - base_scheduler = StepLRScheduler( - optimizer, - step_size=cfg.lr_scheduler.get("step_size", 1000), - gamma=cfg.lr_scheduler.get("gamma", 0.1), - ) - elif scheduler_type == "CosineAnnealingLR": - base_scheduler = CosineAnnealingLR( - optimizer, - t_max=cfg.lr_scheduler.get("t_max", 5000), - eta_min=cfg.lr_scheduler.get("eta_min", 0.0), - warm=cfg.lr_scheduler.get("warm", False), - ) - elif scheduler_type == "WarmupScheduler": - # Create nested base scheduler - base_cfg = cfg.lr_scheduler.base_scheduler - if base_cfg.type == "StepLRScheduler": - nested_scheduler = StepLRScheduler( - optimizer, - step_size=base_cfg.get("step_size", 1000), - gamma=base_cfg.get("gamma", 0.1), - ) - elif base_cfg.type == "CosineAnnealingLR": - nested_scheduler = CosineAnnealingLR( - optimizer, - t_max=base_cfg.get("t_max", 5000), - eta_min=base_cfg.get("eta_min", 0.0), - warm=base_cfg.get("warm", False), - ) - else: - raise ValueError(f"Unknown base scheduler type: {base_cfg.type}") - - return WarmupScheduler( - optimizer, - base_scheduler=nested_scheduler, - warmup_steps=cfg.lr_scheduler.get("warmup_steps", 500), - warmup_start_lr=cfg.lr_scheduler.get("warmup_start_lr", 0.0), - ) - else: - raise ValueError(f"Unknown LR scheduler type: {scheduler_type}") - - return base_scheduler - - -def create_solver(cfg: DictConfig, model, preprocess_hook): - """Create ODE solver from config.""" - solver_type = cfg.solver.type - if solver_type == "euler": - return Euler(model, preprocess_hook=preprocess_hook) - if solver_type == "heun": - return Heun(model, preprocess_hook=preprocess_hook) - if solver_type == "midpoint": - return MidpointSolver(model, preprocess_hook=preprocess_hook) - if solver_type == "rk4": - return RK4(model, preprocess_hook=preprocess_hook) - if solver_type == "ddim": - eta = cfg.solver.get("eta", 0.0) - return DDIM(model, preprocess_hook=preprocess_hook, eta=eta) - raise ValueError(f"Unknown solver type: {solver_type}") - - -def create_model(cfg: DictConfig): - """Create model from config.""" - model_type = cfg.model.type - if model_type == "neural_network": - return NeuralNetwork( - in_dim=cfg.model.input_dim, - time_embed_dim=cfg.model.time_embed_dim, - out_dim=cfg.model.output_dim, - ) - else: - raise ValueError(f"Unknown model type: {model_type}") - - -def generate_model_name(cfg: DictConfig) -> str: - """Generate model name from config.""" - dataset = cfg.dataset.name - model_type = cfg.model.type - scheduler = cfg.scheduler.type.replace("Scheduler", "").lower() - name = f"model_{dataset}_{model_type}_{scheduler}.safetensors" - return name - - -def epoch(x_1, model, path): - """Single training epoch.""" - x_1 = T(x_1.astype("float32")) # pyright: ignore - t = T.rand(x_1.shape[0], 1) * 0.99 # clamping - x_0 = T.randn(*x_1.shape) - x_t, dx_t = path.sample(x_1=x_1, t=t, x_0=x_0) - out = model(x_t, t=t) # pyright: ignore - return out, dx_t - - -@hydra.main(version_base=None, config_path="../configs", config_name="config") -def main(cfg: DictConfig): - """Main training function.""" - print("Configuration:") - print(OmegaConf.to_yaml(cfg)) - - # Set random seed if specified - if cfg.get("seed"): - T.manual_seed(cfg.seed) - - # Generate model name (can be overridden with model_name parameter) - model_name = cfg.get("model_name", generate_model_name(cfg)) - print(f"\nModel will be saved as: {model_name}") - - # Create model, scheduler, and path - model = create_model(cfg) - scheduler = create_scheduler(cfg) - path = AffinePath(scheduler=scheduler) - - # Create optimizer - optim = Adam(get_parameters(model), lr=cfg.optimizer.lr) - - # Create learning rate scheduler - lr_scheduler = create_lr_scheduler(cfg, optim) - - loss_fn = mse - _losses = [] - - # Training loop - pbar = tqdm(range(cfg.training.num_epochs)) - T.training = True - for iter_idx in pbar: - x, _ = make_moons( - n_samples=cfg.dataset.n_samples, - noise=cfg.dataset.noise, - ) - out, dx_t = epoch(x, model, path) - optim.zero_grad() - loss = loss_fn(out, dx_t) - loss.backward() - loss_val = loss.item() - _losses.append(loss_val) - - if iter_idx % cfg.training.log_interval == 0: - desc = f"Loss: {loss_val:.4e}" - if lr_scheduler is not None: - current_lr = lr_scheduler.get_lr() - desc += f" | LR: {current_lr:.6f}" - pbar.set_description_str(desc) - - optim.step() - - # Step learning rate scheduler per iteration - if lr_scheduler is not None: - lr_scheduler.step(iter_idx) - - # Save model after training - if cfg.training.get("save_model", True): - from tinygrad.nn.state import get_state_dict, safe_save - - safe_save(get_state_dict(model), model_name) - print(f"✓ Model saved to: {model_name}") - - # Plot loss curve after training - if cfg.training.get("log_artifacts", True): - output_dir = cfg.get("output_dir", "outputs") - os.makedirs(output_dir, exist_ok=True) - - _ = plt.figure(figsize=(10, 4)) - plt.plot(_losses) - plt.xlabel("Iteration") - plt.ylabel("Loss") - plt.title("Training Loss Over Time") - plt.grid(True) - plt.tight_layout() - - loss_path = os.path.join(output_dir, "loss_curve.png") - plt.savefig(loss_path) - plt.close() - - # Generate samples after training - if cfg.training.get("generate_samples", True): - num_samples = cfg.training.get("n_samples", 100) - x = T.randn(num_samples, 2) - h_step = cfg.training.step_size - time_grid = T.linspace(0, 1, int(1 / h_step)) - - solver = create_solver(cfg, model, preprocess_time_moons) - visualize_moons( - x, - solver=solver, - time_grid=time_grid, - h_step=h_step, - num_plots=cfg.training.get("num_plots", 10), - ) - - -if __name__ == "__main__": - main() diff --git a/experiments/_common.py b/experiments/_common.py new file mode 100644 index 0000000..35d1ef7 --- /dev/null +++ b/experiments/_common.py @@ -0,0 +1,45 @@ +"""Shared helpers for the ODE-solver/schedule research scripts in this directory.""" + +import matplotlib.pyplot as plt +import numpy as np +from tinygrad.nn.state import safe_load + + +def detect_time_embed_dim(model_path: str, in_channels: int) -> int: + """Detect time_embed_dim from saved weights by inspecting enc1.conv.weight shape.""" + state = safe_load(model_path) + key = "enc1.conv.weight" + if key in state: + return int(state[key].shape[1]) - in_channels + return 64 + + +def schedule_grid(N: int, p: float) -> np.ndarray: + """Time grid t_k = 1 - (1 - k/N)^p; p=1 is uniform, p>1 clusters steps near t=1.""" + k = np.arange(N + 1) / N + return 1.0 - (1.0 - k) ** p + + +def normalize_for_plot(x_np: np.ndarray) -> np.ndarray: + out = (x_np - x_np.min()) / (x_np.max() - x_np.min() + 1e-8) + return np.clip(out, 0, 1) + + +def make_grid(x_np: np.ndarray, grid: int = 3) -> np.ndarray: + h, w = x_np.shape[-2:] + canvas = np.ones((grid * h + (grid - 1), grid * w + (grid - 1))) + for i in range(grid): + for j in range(grid): + idx = i * grid + j + canvas[i * (h + 1) : i * (h + 1) + h, j * (w + 1) : j * (w + 1) + w] = x_np[idx, 0] + return canvas + + +def save_grid_figure(x_np: np.ndarray, title: str, path: str, figsize=(4, 4.2)): + fig, ax = plt.subplots(figsize=figsize) + ax.imshow(make_grid(normalize_for_plot(x_np)), cmap="gray") + ax.set_title(title, fontsize=11) + ax.axis("off") + plt.tight_layout() + plt.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) diff --git a/experiments/curvature_profile.py b/experiments/curvature_profile.py new file mode 100644 index 0000000..f87f994 --- /dev/null +++ b/experiments/curvature_profile.py @@ -0,0 +1,294 @@ +""" +Curvature profile for trained flow-matching models. Logs to MLflow. + +At each Euler step, finite-difference-estimate: + - ||v|| velocity magnitude + - ||d_t v|| partial wrt time + - ||J_x v · v_hat|| directional derivative along streamline + - ||J_x v · eta|| Hutchinson-style random-direction proxy + - lte_proxy = ||d_t v|| + ||J_x v · v_hat|| * ||v|| (Euler local truncation / h^2) + +Usage: + uv run experiments/curvature_profile.py --dataset moons + uv run experiments/curvature_profile.py --dataset mnist --num-steps 50 +""" + +import argparse +import os + +import matplotlib.pyplot as plt +import mlflow +import numpy as np +from _common import detect_time_embed_dim, schedule_grid +from tinygrad.tensor import Tensor as T +from tqdm import tqdm + +from tinyflow.nn import NeuralNetwork, UNetTinygrad +from tinyflow.trainer import BaseTrainer +from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist, preprocess_time_moons + +plt.style.use("ggplot") + +mlflow.set_tracking_uri("sqlite:///mlflow.db") + + +DATASET_CONFIGS = { + "moons": { + "shape": (2,), + "kind": "mlp", + "preprocess": preprocess_time_moons, + "default_path": "model_moons_neural_network_linear.safetensors", + "default_batch": 256, + }, + "mnist": { + "shape": (1, 28, 28), + "kind": "unet", + "in_channels": 1, + "preprocess": preprocess_time_mnist, + "default_path": "model_mnist_unet_linear.safetensors", + "default_batch": 32, + }, + "fashion_mnist": { + "shape": (1, 28, 28), + "kind": "unet", + "in_channels": 1, + "preprocess": preprocess_time_mnist, + "default_path": "model_fashion_mnist_unet_linear.safetensors", + "default_batch": 32, + }, + "cifar10": { + "shape": (3, 32, 32), + "kind": "unet", + "in_channels": 3, + "preprocess": preprocess_time_cifar, + "default_path": "model_cifar10_unet_linear.safetensors", + "default_batch": 16, + }, +} + + +def build_model(dataset: str, model_path: str): + cfg = DATASET_CONFIGS[dataset] + if cfg["kind"] == "mlp": + return NeuralNetwork(in_dim=2, time_embed_dim=64, out_dim=2) + in_ch = cfg["in_channels"] + 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 +): + cfg = DATASET_CONFIGS[dataset] + preprocess = cfg["preprocess"] + + model = build_model(dataset, model_path) + BaseTrainer.load_model(model, model_path) + + T.training = False + shape = tuple([batch_size] + list(cfg["shape"])) + x = T.randn(*shape).realize() + h = 1.0 / num_steps + + records = [] + for step in tqdm(range(num_steps), desc="profiling"): + t_val = step * h + t = (T.zeros(1) + t_val).contiguous() + t_in = preprocess(t, x) + + v = model(x, t_in).realize() + v_norm = per_sample_norm_mean(v) + + t_plus_val = min(t_val + eps_t, 1.0) + t_plus = preprocess((T.zeros(1) + t_plus_val).contiguous(), x) + v_tplus = model(x, t_plus).realize() + dvdt_norm = per_sample_norm_mean((v_tplus - v) * (1.0 / max(t_plus_val - t_val, 1e-12))) + + v_hat = normalize_per_sample(v) + v_at_xv = model(x + eps_x * v_hat, t_in).realize() + jvv_norm = per_sample_norm_mean((v_at_xv - v) * (1.0 / eps_x)) + + eta_hat = normalize_per_sample(T.randn(*shape).realize()) + v_at_eta = model(x + eps_x * eta_hat, t_in).realize() + jeta_norm = per_sample_norm_mean((v_at_eta - v) * (1.0 / eps_x)) + + lte_proxy = dvdt_norm + jvv_norm * v_norm + record = { + "t": t_val, + "v": v_norm, + "dvdt": dvdt_norm, + "Jvv": jvv_norm, + "Jeta": jeta_norm, + "lte_proxy": lte_proxy, + } + mlflow.log_metrics(record, step=step) + records.append(record) + + x = (x + h * v).realize() + + return records + + +def make_basic_plot(records, dataset: str): + ts = [r["t"] for r in records] + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + + ax = axes[0] + ax.plot(ts, [r["v"] for r in records], label="‖v‖", linewidth=2) + ax.plot(ts, [r["dvdt"] for r in records], label="‖∂_t v‖", linewidth=2) + ax.plot(ts, [r["Jvv"] for r in records], label="‖J_x v · v̂‖", linewidth=2) + ax.plot(ts, [r["Jeta"] for r in records], label="‖J_x v · η̂‖", linewidth=2) + ax.set_xlabel("t") + ax.set_ylabel("magnitude") + ax.set_title(f"Curvature profile — {dataset}") + ax.legend() + + ax = axes[1] + ax.plot(ts, [r["lte_proxy"] for r in records], label="LTE proxy", linewidth=2, color="C3") + ax.set_xlabel("t") + ax.set_ylabel("LTE proxy / h²") + ax.set_title("Per-step error growth proxy") + ax.legend() + + plt.tight_layout() + return fig + + +def make_story_plot(records, dataset: str, N: int, p: float): + ts = np.array([r["t"] for r in records]) + v = np.array([r["v"] for r in records]) + dvdt = np.array([r["dvdt"] for r in records]) + Jvv = np.array([r["Jvv"] for r in records]) + lte = np.array([r["lte_proxy"] for r in records]) + + fig, axes = plt.subplots( + 2, 1, figsize=(11, 7), gridspec_kw={"height_ratios": [3, 1]}, sharex=True + ) + + ax = axes[0] + ax.plot(ts, v, label="‖v‖ (sample velocity)", linewidth=2.2) + ax.plot(ts, dvdt, label="‖∂_t v‖", linewidth=2.2) + ax.plot(ts, Jvv, label="‖J_x v · v̂‖", linewidth=2.2) + ax.plot(ts, lte, label="LTE proxy", linewidth=2.5, color="C3", linestyle="--") + ax.set_yscale("log") + ax.set_ylabel("magnitude (log scale)") + ax.set_title(f"Magnitude profile along Euler trajectory — {dataset}") + ax.legend(loc="upper left", fontsize=9) + ax.grid(True, which="both", alpha=0.3) + + above = np.where(lte > 2.0 * lte[0])[0] + if above.size > 0: + t_thr = ts[above[0]] + ax.axvspan(t_thr, 1.0, color="orange", alpha=0.12) + ax.text( + (t_thr + 1.0) / 2, + ax.get_ylim()[1] * 0.4, + f"LTE proxy > 2× start\n(t > {t_thr:.2f})", + ha="center", + fontsize=9, + color="C1", + ) + + ax = axes[1] + uniform = schedule_grid(N, 1.0) + backloaded = schedule_grid(N, p) + ax.scatter(uniform, np.full_like(uniform, 1.0), marker="|", s=400, color="C0") + ax.scatter(backloaded, np.full_like(backloaded, 0.0), marker="|", s=400, color="C2") + ax.text(-0.04, 1.0, "uniform", ha="right", va="center", fontsize=10) + ax.text(-0.04, 0.0, f"back-loaded\n(p={p})", ha="right", va="center", fontsize=10) + ax.set_yticks([]) + ax.set_ylim(-0.5, 1.5) + ax.set_xlim(-0.10, 1.02) + ax.set_xlabel("t") + ax.set_title(f"Step allocation comparison (N={N})") + ax.grid(True, axis="x", alpha=0.3) + + plt.tight_layout() + return fig + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", choices=list(DATASET_CONFIGS), required=True) + parser.add_argument("--model-path", default=None) + parser.add_argument("--num-steps", type=int, default=50) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--eps-x", type=float, default=1e-2) + parser.add_argument("--eps-t", type=float, default=1e-3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--schedule-N", + type=int, + default=20, + help="N for the step-allocation strip in the story plot", + ) + parser.add_argument( + "--schedule-p", + type=float, + default=2.5, + help="back-loading exponent for the comparison strip", + ) + parser.add_argument("--experiment", default="curvature_profiling") + args = parser.parse_args() + + T.manual_seed(args.seed) + + cfg = DATASET_CONFIGS[args.dataset] + model_path = args.model_path or cfg["default_path"] + batch_size = args.batch_size or cfg["default_batch"] + + if not os.path.exists(model_path): + raise FileNotFoundError(f"Model file not found: {model_path}") + + mlflow.set_experiment(args.experiment) + with mlflow.start_run(run_name=f"{args.dataset}_N{args.num_steps}"): + mlflow.set_tags({"dataset": args.dataset, "kind": "curvature"}) + mlflow.log_params( + { + "dataset": args.dataset, + "model_path": model_path, + "num_steps": args.num_steps, + "batch_size": batch_size, + "eps_x": args.eps_x, + "eps_t": args.eps_t, + "seed": args.seed, + } + ) + + records = profile( + args.dataset, model_path, args.num_steps, batch_size, args.eps_x, args.eps_t + ) + + basic = make_basic_plot(records, args.dataset) + mlflow.log_figure(basic, f"curvature_{args.dataset}.png") + plt.close(basic) + + story = make_story_plot(records, args.dataset, args.schedule_N, args.schedule_p) + mlflow.log_figure(story, f"story_{args.dataset}.png") + plt.close(story) + + lte = np.array([r["lte_proxy"] for r in records]) + ratio = float(lte.max() / max(lte.min(), 1e-12)) + mlflow.log_metrics( + { + "summary_lte_max_over_min": ratio, + "summary_lte_max": float(lte.max()), + "summary_lte_min": float(lte.min()), + } + ) + print(f"done. ratio max/min(lte) = {ratio:.2f}") + + +if __name__ == "__main__": + main() diff --git a/experiments/schedule_compare.py b/experiments/schedule_compare.py new file mode 100644 index 0000000..0322dc5 --- /dev/null +++ b/experiments/schedule_compare.py @@ -0,0 +1,138 @@ +""" +Compare uniform vs back-loaded time schedules on a trained MNIST model. + +Both schedules sample the same Euler ODE; only the time grid differs: + - uniform: t_k = k / N + - back-loaded: t_k = 1 - (1 - k/N)^p (clusters steps near t=1) + +Same starting noise (fixed seed) for every (schedule, N) combo so any +visual difference is attributable to the schedule, not the seed. + +Each (schedule, N) is one MLflow run; artifact is a 9-sample grid. +A combined side-by-side figure is logged to a parent run. + +Usage: + uv run experiments/schedule_compare.py + uv run experiments/schedule_compare.py --steps 10 20 50 --p 2.5 +""" + +import argparse +import os + +import matplotlib.pyplot as plt +import mlflow +import numpy as np +from _common import ( + detect_time_embed_dim, + make_grid, + normalize_for_plot, + save_grid_figure, + schedule_grid, +) +from tinygrad.tensor import Tensor as T +from tqdm import tqdm + +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 euler_integrate(model, x0, ts: np.ndarray, preprocess): + x = x0 + for i in tqdm(range(len(ts) - 1), desc=f"integrate N={len(ts) - 1}", leave=False): + t = (T.zeros(1) + float(ts[i])).contiguous() + h = float(ts[i + 1] - ts[i]) + v = model(x, preprocess(t, x)) + x = (x + h * v).realize() + return x + + +def make_combined_figure(results, p: float, path: str): + Ns = sorted({N for (_, N) in results}) + fig, axes = plt.subplots(len(Ns), 2, figsize=(8, 4 * len(Ns))) + if len(Ns) == 1: + axes = axes[None, :] + for row, N in enumerate(Ns): + for col, (label, sched_p) in enumerate([("uniform", 1.0), (f"back-loaded p={p}", p)]): + x_np = results[(sched_p, N)] + axes[row, col].imshow(make_grid(normalize_for_plot(x_np)), cmap="gray") + axes[row, col].set_title(f"{label}, N={N} (NFE={N})", fontsize=11) + axes[row, col].axis("off") + plt.tight_layout() + plt.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", default="model_mnist_unet_linear.safetensors") + parser.add_argument("--steps", type=int, nargs="+", default=[10, 20, 50]) + parser.add_argument( + "--p", + type=float, + default=2.5, + help="back-loading exponent; 1=uniform, >1 clusters near t=1", + ) + parser.add_argument("--n-samples", type=int, default=9) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--experiment", default="schedule_compare") + 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) + + schedules = [("uniform", 1.0), ("back-loaded", args.p)] + results = {} + + with mlflow.start_run(run_name=f"compare_p{args.p}") as parent: + mlflow.set_tags({"kind": "schedule_compare", "dataset": "mnist"}) + mlflow.log_params( + { + "model_path": args.model_path, + "p": args.p, + "n_samples": args.n_samples, + "seed": args.seed, + "steps": ",".join(str(s) for s in args.steps), + } + ) + + for N in args.steps: + for label, sched_p in schedules: + T.manual_seed(args.seed) + x0 = T.randn(args.n_samples, 1, 28, 28).realize() + + ts = schedule_grid(N, sched_p) + with mlflow.start_run(run_name=f"{label}_N{N}", nested=True) as child: + mlflow.set_tags({"schedule": label, "kind": "schedule_compare"}) + mlflow.log_params({"schedule": label, "p": sched_p, "N": N, "NFE": N}) + + x_final = euler_integrate(model, x0, ts, preprocess_time_mnist) + x_np = x_final.numpy() + results[(sched_p, N)] = x_np + + grid_path = f"/tmp/grid_{label}_N{N}.png" + save_grid_figure(x_np, f"{label}, N={N}", grid_path, figsize=(4, 4.4)) + mlflow.log_artifact(grid_path) + print(f"logged child run {child.info.run_id}: {label} N={N}") + + combined_path = "/tmp/schedule_compare_mnist.png" + make_combined_figure(results, args.p, combined_path) + mlflow.log_artifact(combined_path) + print(f"parent run: {parent.info.run_id}") + print(f"combined figure: {combined_path}") + + +if __name__ == "__main__": + main() diff --git a/experiments/solver_compare.py b/experiments/solver_compare.py new file mode 100644 index 0000000..e63a4cc --- /dev/null +++ b/experiments/solver_compare.py @@ -0,0 +1,183 @@ +""" +Compare ODE solvers and time schedules at fixed NFE budgets on MNIST. + +Each cell uses the same starting noise (fixed seed) so visual differences +trace to the solver/schedule, not the seed. + +NFE accounting (model forward calls): + Euler : 1 eval/step -> N = NFE + Heun : 2 evals/step -> N = NFE / 2 + RK4 : 4 evals/step -> N = NFE / 4 + +Configs at each NFE budget: + Euler uniform t_k = k/N + Euler back-loaded (p=2.5) t_k = 1 - (1 - k/N)^p + Heun uniform + RK4 uniform + +Usage: + uv run experiments/solver_compare.py + uv run experiments/solver_compare.py --nfes 8 12 20 40 --p 2.5 +""" + +import argparse +import os + +import matplotlib.pyplot as plt +import mlflow +import numpy as np +from _common import ( + detect_time_embed_dim, + make_grid, + normalize_for_plot, + save_grid_figure, + schedule_grid, +) +from tinygrad.tensor import Tensor as T + +from tinyflow.nn import UNetTinygrad +from tinyflow.solver import RK4, Euler, Heun +from tinyflow.trainer import BaseTrainer +from tinyflow.utils import preprocess_time_mnist + +plt.style.use("ggplot") +mlflow.set_tracking_uri("sqlite:///mlflow.db") + + +CONFIGS = [ + ("euler_uniform", "Euler uniform", "euler", 1.0, 1), + ("euler_backloaded", "Euler back-loaded", "euler", None, 1), + ("heun_uniform", "Heun uniform", "heun", 1.0, 2), + ("rk4_uniform", "RK4 uniform", "rk4", 1.0, 4), +] + + +def make_solver(name: str, model): + if name == "euler": + return Euler(model, preprocess_hook=preprocess_time_mnist) + if name == "heun": + return Heun(model, preprocess_hook=preprocess_time_mnist) + if name == "rk4": + return RK4(model, preprocess_hook=preprocess_time_mnist) + raise ValueError(name) + + +def integrate(solver, x0, ts: np.ndarray): + x = x0 + for i in range(len(ts) - 1): + t = (T.zeros(1) + float(ts[i])).contiguous() + h = float(ts[i + 1] - ts[i]) + x = solver.sample(h, t, x).realize() + return x + + +def make_combined_figure(results, nfes, path: str): + n_rows = len(nfes) + n_cols = len(CONFIGS) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.2 * n_cols, 3.5 * n_rows)) + if n_rows == 1: + axes = axes[None, :] + for r, nfe in enumerate(nfes): + for c, (key, label, _, _, _) in enumerate(CONFIGS): + ax = axes[r, c] + x_np = results.get((key, nfe)) + if x_np is None: + ax.text( + 0.5, + 0.5, + "skipped\n(N<1)", + ha="center", + va="center", + fontsize=12, + color="gray", + transform=ax.transAxes, + ) + ax.set_xticks([]) + ax.set_yticks([]) + else: + ax.imshow(make_grid(normalize_for_plot(x_np)), cmap="gray") + ax.axis("off") + ax.set_title(f"{label}\nNFE={nfe}", fontsize=10) + plt.tight_layout() + plt.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", default="model_mnist_unet_linear.safetensors") + parser.add_argument("--nfes", type=int, nargs="+", default=[8, 12, 20, 40]) + parser.add_argument("--p", type=float, default=2.5) + parser.add_argument("--n-samples", type=int, default=9) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--experiment", default="solver_compare") + 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) + results = {} + + with mlflow.start_run(run_name=f"compare_p{args.p}") as parent: + mlflow.set_tags({"kind": "solver_compare", "dataset": "mnist"}) + mlflow.log_params( + { + "model_path": args.model_path, + "p": args.p, + "n_samples": args.n_samples, + "seed": args.seed, + "nfes": ",".join(str(s) for s in args.nfes), + } + ) + + for nfe in args.nfes: + for key, label, solver_name, default_p, evals_per_step in CONFIGS: + p_val = default_p if default_p is not None else args.p + N = nfe // evals_per_step + if N < 1: + print(f" skip {key} NFE={nfe}: would need N={N}") + continue + + T.manual_seed(args.seed) + x0 = T.randn(args.n_samples, 1, 28, 28).realize() + solver = make_solver(solver_name, model) + ts = schedule_grid(N, p_val) + + with mlflow.start_run(run_name=f"{key}_NFE{nfe}", nested=True) as child: + mlflow.set_tags({"config": key, "kind": "solver_compare"}) + mlflow.log_params( + { + "config": key, + "solver": solver_name, + "p": p_val, + "N": N, + "NFE": nfe, + } + ) + + x_final = integrate(solver, x0, ts) + x_np = x_final.numpy() + results[(key, nfe)] = x_np + + grid_path = f"/tmp/grid_{key}_NFE{nfe}.png" + save_grid_figure(x_np, f"{label} NFE={nfe}", grid_path) + mlflow.log_artifact(grid_path) + + print(f" {key} NFE={nfe} N={N} -> run {child.info.run_id}") + + combined_path = "/tmp/solver_compare_mnist.png" + make_combined_figure(results, args.nfes, combined_path) + mlflow.log_artifact(combined_path) + print(f"parent run: {parent.info.run_id}") + print(f"combined: {combined_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/_common.py b/scripts/_common.py new file mode 100644 index 0000000..fdf6d11 --- /dev/null +++ b/scripts/_common.py @@ -0,0 +1,128 @@ +"""Shared Hydra config factories used by scripts/train.py and scripts/generate.py.""" + +from omegaconf import DictConfig +from tinygrad.nn.state import safe_load + +from tinyflow.nn_utils.lr_scheduler import ( + CosineAnnealingLR, + StepLRScheduler, + WarmupScheduler, +) +from tinyflow.path.scheduler import ( + CosineScheduler, + LinearScheduler, + LinearVarPresScheduler, + PolynomialScheduler, +) +from tinyflow.solver import DDIM, RK4, Euler, Heun, MidpointSolver +from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist, preprocess_time_moons + + +def create_scheduler(cfg: DictConfig): + """Create flow-matching path scheduler from config.""" + scheduler_type = cfg.scheduler.type + if scheduler_type == "LinearScheduler": + return LinearScheduler() + if scheduler_type == "CosineScheduler": + return CosineScheduler() + if scheduler_type == "LinearVarPresScheduler": + return LinearVarPresScheduler() + if scheduler_type == "PolynomialScheduler": + degree = cfg.scheduler.get("degree", 2) + return PolynomialScheduler(degree) + raise ValueError(f"Unknown scheduler type: {scheduler_type}") + + +def create_lr_scheduler(cfg: DictConfig, optimizer): + """Create learning rate scheduler from config.""" + if not cfg.get("lr_scheduler"): + return None + + scheduler_type = cfg.lr_scheduler.type + + # Create base scheduler + if scheduler_type == "NullLRScheduler": + return None + elif scheduler_type == "StepLRScheduler": + base_scheduler = StepLRScheduler( + optimizer, + step_size=cfg.lr_scheduler.get("step_size", 1000), + gamma=cfg.lr_scheduler.get("gamma", 0.1), + ) + elif scheduler_type == "CosineAnnealingLR": + base_scheduler = CosineAnnealingLR( + optimizer, + t_max=cfg.lr_scheduler.get("t_max", 5000), + eta_min=cfg.lr_scheduler.get("eta_min", 0.0), + warm=cfg.lr_scheduler.get("warm", False), + ) + elif scheduler_type == "WarmupScheduler": + # Create nested base scheduler + base_cfg = cfg.lr_scheduler.base_scheduler + if base_cfg.type == "StepLRScheduler": + nested_scheduler = StepLRScheduler( + optimizer, + step_size=base_cfg.get("step_size", 1000), + gamma=base_cfg.get("gamma", 0.1), + ) + elif base_cfg.type == "CosineAnnealingLR": + nested_scheduler = CosineAnnealingLR( + optimizer, + t_max=base_cfg.get("t_max", 5000), + eta_min=base_cfg.get("eta_min", 0.0), + warm=base_cfg.get("warm", False), + ) + else: + raise ValueError(f"Unknown base scheduler type: {base_cfg.type}") + + return WarmupScheduler( + optimizer, + base_scheduler=nested_scheduler, + warmup_steps=cfg.lr_scheduler.get("warmup_steps", 500), + warmup_start_lr=cfg.lr_scheduler.get("warmup_start_lr", 0.0), + ) + else: + raise ValueError(f"Unknown LR scheduler type: {scheduler_type}") + + return base_scheduler + + +def create_solver(cfg: DictConfig, model, preprocess_hook): + """Create ODE solver from config.""" + solver_type = cfg.solver.type + if solver_type == "euler": + return Euler(model, preprocess_hook=preprocess_hook) + if solver_type == "heun": + return Heun(model, preprocess_hook=preprocess_hook) + if solver_type == "midpoint": + return MidpointSolver(model, preprocess_hook=preprocess_hook) + if solver_type == "rk4": + return RK4(model, preprocess_hook=preprocess_hook) + if solver_type == "ddim": + eta = cfg.solver.get("eta", 0.0) + return DDIM(model, preprocess_hook=preprocess_hook, eta=eta) + raise ValueError(f"Unknown solver type: {solver_type}") + + +def get_preprocess_hook(cfg: DictConfig): + """Get the time-preprocessing hook for the configured dataset.""" + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + if dataset_type == "moons": + return preprocess_time_moons + if dataset_type in ["mnist", "fashion_mnist"]: + return preprocess_time_mnist + if dataset_type == "cifar10": + return preprocess_time_cifar + raise ValueError(f"Unknown dataset type: {dataset_type}") + + +def detect_time_embed_dim(model_path: str, in_channels: int) -> int: + """Detect time_embed_dim from saved weights by inspecting enc1.conv.weight shape.""" + try: + state = safe_load(model_path) + key = "enc1.conv.weight" + if key in state: + return int(state[key].shape[1]) - in_channels + except Exception: + pass + return 64 # default diff --git a/scripts/download_datasets.py b/scripts/download_datasets.py index d21a1cd..95d8138 100755 --- a/scripts/download_datasets.py +++ b/scripts/download_datasets.py @@ -10,8 +10,6 @@ import argparse import gzip -import os -import pickle import shutil import tarfile import urllib.request diff --git a/examples/generate_images.py b/scripts/generate.py similarity index 63% rename from examples/generate_images.py rename to scripts/generate.py index 8eb9006..0df2ba0 100644 --- a/examples/generate_images.py +++ b/scripts/generate.py @@ -1,53 +1,52 @@ """ -Generate samples from a trained flow matching model for image datasets (MNIST, Fashion MNIST, CIFAR-10). +Unified generation entry point (Hydra configuration). -Usage: - # Generate static grid for MNIST - uv run examples/generate_images.py generation.model_path=model_mnist_unet_linear.safetensors - - # Generate for Fashion MNIST - uv run examples/generate_images.py dataset=fashion_mnist generation.model_path=model_fashion_mnist_unet_linear.safetensors - - # Generate for CIFAR-10 - uv run examples/generate_images.py dataset=cifar10 generation.model_path=model_cifar10_unet_linear.safetensors - - # Generate with animation - uv run examples/generate_images.py generation.model_path=model_mnist_unet_linear.safetensors --animated +Loads a trained flow-matching model and integrates the learned velocity field +from noise (t=0) to data (t=1). Supports the 2D moons dataset and image +datasets (MNIST, Fashion MNIST, CIFAR-10). - # Generate with class predictions (requires trained FID classifier) - uv run examples/generate_images.py generation.model_path=model.safetensors generation.show_predictions=true +Note: the scheduler/path used during training is not needed here - the model +has already learned the velocity field. - # Generate with FID computation (requires trained FID classifier) - uv run examples/generate_images.py generation.model_path=model.safetensors generation.compute_fid=true +Usage: + # Image datasets (default config base: MNIST) + uv run scripts/generate.py generation.model_path=model_mnist_unet_linear.safetensors + uv run scripts/generate.py dataset=fashion_mnist generation.model_path=model_fashion_mnist_unet_linear.safetensors + uv run scripts/generate.py dataset=cifar10 generation.model_path=model_cifar10_unet_linear.safetensors - # Generate with both predictions and FID - uv run examples/generate_images.py generation.model_path=model.safetensors generation.show_predictions=true generation.compute_fid=true + # Animated GIF, class predictions, FID + uv run scripts/generate.py generation.model_path=model.safetensors --animated + uv run scripts/generate.py generation.model_path=model.safetensors generation.show_predictions=true generation.compute_fid=true # Override grid size and steps - uv run examples/generate_images.py generation.model_path=model.safetensors generation.grid_size=3 generation.num_steps=50 + uv run scripts/generate.py generation.model_path=model.safetensors generation.grid_size=3 generation.num_steps=50 + + # Moons dataset - switch the base config with --config-name + uv run scripts/generate.py --config-name=generate_moons_config generation.model_path=model_moons_neural_network_linear.safetensors + uv run scripts/generate.py --config-name=generate_moons_config generation.model_path=model.safetensors --animated """ import os +import sys import hydra import matplotlib.pyplot as plt import numpy as np +from _common import create_solver, detect_time_embed_dim from omegaconf import DictConfig, OmegaConf from PIL import Image from tinygrad import TinyJit -from tinygrad.nn.state import safe_load from tinygrad.tensor import Tensor as T from tqdm import tqdm from tinyflow.metrics import get_feature_extractor -from tinyflow.nn import UNetTinygrad -from tinyflow.solver import DDIM, RK4, Euler, Heun, MidpointSolver +from tinyflow.nn import NeuralNetwork, UNetTinygrad from tinyflow.trainer import BaseTrainer -from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist +from tinyflow.utils import preprocess_time_cifar, preprocess_time_mnist, preprocess_time_moons plt.style.use("ggplot") -# Class names for datasets +# Class names for image datasets CLASS_NAMES = { "mnist": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"], "fashion_mnist": [ @@ -77,49 +76,31 @@ } -def _detect_time_embed_dim(model_path: str, in_channels: int) -> int: - """Detect time_embed_dim from saved weights by inspecting enc1.conv.weight shape.""" - try: - state = safe_load(model_path) - key = "enc1.conv.weight" - if key in state: - return int(state[key].shape[1]) - in_channels - except Exception: - pass - return 64 # default - - def create_model(cfg: DictConfig): - """Create model from config.""" + """Create model from config, sized to match the checkpoint being loaded.""" model_type = cfg.model.type + if model_type == "neural_network": + return NeuralNetwork( + in_dim=cfg.model.input_dim, + time_embed_dim=cfg.model.time_embed_dim, + out_dim=cfg.model.output_dim, + ) + dataset_type = cfg.dataset.get("type", cfg.dataset.name) model_path = cfg.generation.model_path - if model_type == "unet": if "mnist" in dataset_type: - time_embed_dim = _detect_time_embed_dim(model_path, in_channels=1) + time_embed_dim = detect_time_embed_dim(model_path, in_channels=1) return UNetTinygrad(time_embed_dim=time_embed_dim) if "cifar" in dataset_type: - time_embed_dim = _detect_time_embed_dim(model_path, in_channels=3) + time_embed_dim = detect_time_embed_dim(model_path, in_channels=3) return UNetTinygrad(3, 3, time_embed_dim=time_embed_dim) raise ValueError(f"Unknown model type: {model_type}") -def create_solver(cfg: DictConfig, model, preprocess_hook): - """Create ODE solver from config.""" - solver_type = cfg.solver.type - if solver_type == "euler": - return Euler(model, preprocess_hook=preprocess_hook) - if solver_type == "heun": - return Heun(model, preprocess_hook=preprocess_hook) - if solver_type == "midpoint": - return MidpointSolver(model, preprocess_hook=preprocess_hook) - if solver_type == "rk4": - return RK4(model, preprocess_hook=preprocess_hook) - if solver_type == "ddim": - eta = cfg.solver.get("eta", 0.0) - return DDIM(model, preprocess_hook=preprocess_hook, eta=eta) - raise ValueError(f"Unknown solver type: {solver_type}") +# -------------------------------------------------------------------------- +# Image datasets (MNIST, Fashion MNIST, CIFAR-10) +# -------------------------------------------------------------------------- def get_dataset_config(cfg: DictConfig): @@ -269,7 +250,7 @@ def classify_images(images_np: np.ndarray, dataset_name: str): return None -def generate_static_grid(cfg: DictConfig, model, solver, dataset_config): +def generate_image_grid(cfg: DictConfig, model, solver, dataset_config): """Generate a static grid of samples.""" grid_size = cfg.generation.grid_size num_steps = cfg.generation.num_steps @@ -362,7 +343,7 @@ def jit_step(h, t, x): plt.close() -def generate_animation(cfg: DictConfig, model, solver, dataset_config): +def generate_image_animation(cfg: DictConfig, model, solver, dataset_config): """Generate an animated GIF of the generation process.""" grid_size = cfg.generation.grid_size num_steps = cfg.generation.num_steps @@ -471,19 +452,173 @@ def jit_step(h, t, x): print(f" Duration: {len(frames) / fps:.1f} seconds") -def main_impl(cfg: DictConfig): - """Main generation function.""" - animated = cfg.generation.get("animated", False) +def run_image_generation(cfg: DictConfig, animated: bool): + """Load an image-dataset model and generate a grid or animation.""" + model_path = cfg.generation.model_path + if not os.path.exists(model_path): + raise FileNotFoundError(f"Model file not found: {model_path}") - print("Configuration:") - print(OmegaConf.to_yaml(cfg)) - print(f"\nAnimated: {animated}") + print(f"\nLoading model from {model_path}...") + model = create_model(cfg) + BaseTrainer.load_model(model, model_path) - # Set random seed - if cfg.get("seed"): - T.manual_seed(cfg.seed) + # Note: Scheduler/path not needed at generation time - the model has already + # learned the velocity field during training + dataset_config = get_dataset_config(cfg) + solver = create_solver(cfg, model, dataset_config["preprocess_hook"]) + + if animated: + generate_image_animation(cfg, model, solver, dataset_config) + else: + generate_image_grid(cfg, model, solver, dataset_config) + + +# -------------------------------------------------------------------------- +# Moons (2D toy dataset) +# -------------------------------------------------------------------------- + + +def generate_moons_grid(cfg: DictConfig, solver): + """Generate static visualization with multiple time snapshots.""" + n_samples = cfg.generation.n_samples + num_steps = cfg.generation.num_steps + step_size = cfg.generation.step_size + num_plots = cfg.generation.num_plots + + print(f"Generating {n_samples} samples with {num_plots} snapshots...") + + # Initialize samples + T.training = False + x = T.randn(n_samples, 2) + time_grid = T.linspace(0, 1, num_steps) + sample_every = time_grid.shape[0] // num_plots + + # Store snapshots to visualize + snapshots = [] + snapshot_times = [] + + # JIT compile the solver step for better performance + @TinyJit + def jit_step(step_size, t, x): + return solver.sample(step_size, t, x) + + # Generate all samples first + for idx in tqdm(range(int(time_grid.shape[0])), desc="Generating samples"): + t = time_grid[idx].contiguous() + x = jit_step(step_size, t, x) + + # Store reference for visualization + if (idx + 1) % sample_every == 0: + snapshots.append(x) + snapshot_times.append(t) + + # Visualize all snapshots + _, ax = plt.subplots(1, num_plots, figsize=(30, 4), sharex=True, sharey=True) + for i, (snapshot, t) in enumerate(zip(snapshots, snapshot_times, strict=False)): + x_np = snapshot.numpy() + ax[i].scatter(x_np[:, 0], x_np[:, 1], s=20, alpha=0.6) + ax[i].set_title(f"t={t.numpy():.2f}") + ax[i].set_xlim(-3, 3) + ax[i].set_ylim(-3, 3) + ax[i].grid(True, alpha=0.3) + + plt.tight_layout() + + # Save + os.makedirs(cfg.generation.output_dir, exist_ok=True) + output_file = os.path.join(cfg.generation.output_dir, "moons_generation.png") + plt.savefig(output_file, dpi=150, bbox_inches="tight") + print(f"Saved to {output_file}") + + if cfg.generation.get("show_plot", False): + plt.show() + else: + plt.close() - # Load model + +def generate_moons_animation(cfg: DictConfig, solver): + """Generate an animated GIF of the generation process.""" + n_samples = cfg.generation.n_samples + num_steps = cfg.generation.num_steps + step_size = cfg.generation.step_size + num_plots = cfg.generation.num_plots + fps = cfg.generation.fps + + print(f"Generating {n_samples} samples with {num_plots} frames...") + + # Initialize samples + T.training = False + x = T.randn(n_samples, 2) + time_grid = T.linspace(0, 1, num_steps) + sample_every = time_grid.shape[0] // num_plots + + # Store snapshots to visualize + snapshots = [] + snapshot_times = [] + + # JIT compile the solver step for better performance + @TinyJit + def jit_step(step_size, t, x): + return solver.sample(step_size, t, x) + + # Generate all samples first + for idx in tqdm(range(int(time_grid.shape[0])), desc="Generating animation"): + t = time_grid[idx].contiguous() + x = jit_step(step_size, t, x) + + # Store reference for visualization + if (idx + 1) % sample_every == 0: + snapshots.append(x) + snapshot_times.append(t) + + # Create frames + frames = [] + for snapshot, t in tqdm( + zip(snapshots, snapshot_times, strict=False), desc="Creating frames", total=len(snapshots) + ): + x_np = snapshot.numpy() + + fig, ax = plt.subplots(figsize=(8, 8)) + ax.scatter(x_np[:, 0], x_np[:, 1], s=20, alpha=0.6) + ax.set_title(f"t = {t.numpy():.2f}", fontsize=16) + ax.set_xlim(-3, 3) + ax.set_ylim(-3, 3) + ax.grid(True, alpha=0.3) + ax.set_aspect("equal") + + # Convert plot to image + fig.canvas.draw() + # Use buffer_rgba() for compatibility across backends + buf = fig.canvas.buffer_rgba() + frame = np.asarray(buf) + # Convert RGBA to RGB + frame = frame[:, :, :3] + frames.append(Image.fromarray(frame)) + plt.close(fig) + + # Save as GIF + os.makedirs(cfg.generation.output_dir, exist_ok=True) + output_file = os.path.join(cfg.generation.output_dir, "moons_generation.gif") + + # Add final frame multiple times to pause at the end + frames.extend([frames[-1]] * fps) + + frames[0].save( + output_file, + save_all=True, + append_images=frames[1:], + duration=1000 // fps, + loop=0, + ) + + print(f"Animation saved to {output_file}") + print(f" Frames: {len(frames)}") + print(f" FPS: {fps}") + print(f" Duration: {len(frames) / fps:.1f} seconds") + + +def run_moons_generation(cfg: DictConfig, animated: bool): + """Load a moons model and generate a static plot or animation.""" model_path = cfg.generation.model_path if not os.path.exists(model_path): raise FileNotFoundError(f"Model file not found: {model_path}") @@ -492,30 +627,42 @@ def main_impl(cfg: DictConfig): model = create_model(cfg) BaseTrainer.load_model(model, model_path) - # Get dataset configuration # Note: Scheduler/path not needed at generation time - the model has already # learned the velocity field during training - dataset_config = get_dataset_config(cfg) - - # Create solver - solver = create_solver(cfg, model, dataset_config["preprocess_hook"]) + solver = create_solver(cfg, model, preprocess_time_moons) - # Generate samples if animated: - generate_animation(cfg, model, solver, dataset_config) + generate_moons_animation(cfg, solver) + else: + generate_moons_grid(cfg, solver) + + +def main_impl(cfg: DictConfig): + """Resolve config, then dispatch to the moons or image generation path.""" + animated = cfg.generation.get("animated", False) + + print("Configuration:") + print(OmegaConf.to_yaml(cfg)) + print(f"\nAnimated: {animated}") + + # Set random seed + if cfg.get("seed"): + T.manual_seed(cfg.seed) + + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + if dataset_type == "moons": + run_moons_generation(cfg, animated) else: - generate_static_grid(cfg, model, solver, dataset_config) + run_image_generation(cfg, animated) @hydra.main(version_base=None, config_path="../configs", config_name="generate_config") def main(cfg: DictConfig): - """Hydra entry point.""" + """Hydra entry point. Use --config-name=generate_moons_config for the moons dataset.""" main_impl(cfg) if __name__ == "__main__": - import sys - # Check for --animated flag and remove it before Hydra processes args animated = "--animated" in sys.argv if animated: diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..150c933 --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,268 @@ +""" +Unified training entry point (Hydra configuration + MLflow tracking). + +Supports the 2D moons toy dataset (MLP) and image datasets - MNIST, Fashion +MNIST, CIFAR-10 (UNet) - selected via `dataset=...` / `model=...` overrides. + +Usage: + # Moons dataset (fast, for testing) + uv run scripts/train.py dataset=moons model=neural_network training=moons_default scheduler=linear + + # MNIST training with default config + uv run scripts/train.py + + # Fashion MNIST / CIFAR-10 + uv run scripts/train.py dataset=fashion_mnist + uv run scripts/train.py dataset=cifar10 model=unet + + # Override specific parameters + uv run scripts/train.py scheduler=cosine optimizer.lr=0.01 + + # Compare multiple schedulers (multirun) + uv run scripts/train.py -m scheduler=linear,cosine,polynomial + + # Use a pre-configured experiment + uv run scripts/train.py +experiment=quick_test +""" + +import os + +import hydra +import matplotlib.pyplot as plt +import mlflow +from _common import create_lr_scheduler, create_scheduler, create_solver, get_preprocess_hook +from omegaconf import DictConfig, OmegaConf +from sklearn.datasets import make_moons +from tinygrad.nn.optim import Adam +from tinygrad.nn.state import get_parameters, get_state_dict, safe_save +from tinygrad.tensor import Tensor as T +from tqdm.auto import tqdm + +from tinyflow.dataloader import CIFAR10Loader, FashionMNISTLoader, MNISTLoader +from tinyflow.losses import mse +from tinyflow.nn import NeuralNetwork, UNetCIFAR10, UNetCIFAR10Large, UNetMNIST +from tinyflow.path import AffinePath +from tinyflow.trainer import CIFAR10Trainer, MNISTTrainer +from tinyflow.utils import visualize_moons + +plt.style.use("ggplot") + +mlflow.set_tracking_uri("sqlite:///mlflow.db") +mlflow.set_experiment("flow_matching") + + +def create_model(cfg: DictConfig): + """Create model from config.""" + model_type = cfg.model.type + if model_type == "neural_network": + return NeuralNetwork( + in_dim=cfg.model.input_dim, + time_embed_dim=cfg.model.time_embed_dim, + out_dim=cfg.model.output_dim, + ) + + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + if model_type == "unet": + if dataset_type in ["mnist", "fashion_mnist"]: + return UNetMNIST(in_channels=1, out_channels=1) + elif dataset_type == "cifar10": + return UNetCIFAR10(in_channels=3, out_channels=3) + elif model_type == "unet_large": + if dataset_type == "cifar10": + return UNetCIFAR10Large(in_channels=3, out_channels=3) + else: + raise ValueError(f"unet_large only supports cifar10, got {dataset_type}") + + raise ValueError(f"Unknown model type: {model_type} with dataset type: {dataset_type}") + + +def create_dataloader(cfg: DictConfig): + """Create dataloader from config (image datasets only).""" + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + batch_size = cfg.dataset.get("batch_size", 32) + flatten = cfg.dataset.get("flatten", False) + shuffle = cfg.dataset.get("shuffle", True) + + if dataset_type == "mnist": + return MNISTLoader(flatten=flatten, batch_size=batch_size, shuffle=shuffle) + elif dataset_type == "fashion_mnist": + path = cfg.dataset.get("path", "dataset/fashion_mnist") + train = cfg.dataset.get("train", True) + return FashionMNISTLoader( + path=path, flatten=flatten, batch_size=batch_size, shuffle=shuffle, train=train + ) + elif dataset_type == "cifar10": + path = cfg.dataset.get("path", "dataset/cifar10/cifar-10-batches-py") + cache = cfg.dataset.get("cache", True) + normalize = cfg.dataset.get("normalize", True) + train = cfg.dataset.get("train", True) + return CIFAR10Loader( + path=path, + batch_size=batch_size, + shuffle=shuffle, + train=train, + cache=cache, + normalize=normalize, + ) + else: + raise ValueError(f"Unknown dataset type: {dataset_type}") + + +def create_trainer(cfg: DictConfig, model, dataloader, optim, path, lr_scheduler=None): + """Create trainer from config (image datasets only).""" + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + num_epochs = cfg.training.num_epochs + log_interval = cfg.training.log_interval + gradient_accumulation_steps = cfg.training.get("gradient_accumulation_steps", 1) + + if dataset_type in ["mnist", "fashion_mnist"]: + trainer_cls = MNISTTrainer + elif dataset_type == "cifar10": + trainer_cls = CIFAR10Trainer + else: + raise ValueError(f"Unknown dataset type: {dataset_type}") + + return trainer_cls( + model=model, + dataloader=dataloader, + optim=optim, + loss_fn=mse, + path=path, + num_epochs=num_epochs, + log_interval=log_interval, + lr_scheduler=lr_scheduler, + gradient_accumulation_steps=gradient_accumulation_steps, + ) + + +def generate_model_name(cfg: DictConfig) -> str: + """Generate model name from config.""" + dataset = cfg.dataset.get("type", cfg.dataset.name) + model_type = cfg.model.type + scheduler = cfg.scheduler.type.replace("Scheduler", "").lower() + return f"model_{dataset}_{model_type}_{scheduler}.safetensors" + + +def train_moons(cfg: DictConfig, model_name: str): + """Training loop for the 2D moons toy dataset.""" + model = create_model(cfg) + scheduler = create_scheduler(cfg) + path = AffinePath(scheduler=scheduler) + + optim = Adam(get_parameters(model), lr=cfg.optimizer.lr) + lr_scheduler = create_lr_scheduler(cfg, optim) + + loss_fn = mse + losses = [] + + pbar = tqdm(range(cfg.training.num_epochs)) + T.training = True + for iter_idx in pbar: + x, _ = make_moons(n_samples=cfg.dataset.n_samples, noise=cfg.dataset.noise) + x_1 = T(x.astype("float32")) # pyright: ignore + t = T.rand(x_1.shape[0], 1) * 0.99 # clamping + x_0 = T.randn(*x_1.shape) + x_t, dx_t = path.sample(x_1=x_1, t=t, x_0=x_0) + out = model(x_t, t=t) # pyright: ignore + + optim.zero_grad() + loss = loss_fn(out, dx_t) + loss.backward() + loss_val = loss.item() + losses.append(loss_val) + + if iter_idx % cfg.training.log_interval == 0: + desc = f"Loss: {loss_val:.4e}" + if lr_scheduler is not None: + desc += f" | LR: {lr_scheduler.get_lr():.6f}" + pbar.set_description_str(desc) + + optim.step() + + if lr_scheduler is not None: + lr_scheduler.step(iter_idx) + + if cfg.training.get("save_model", True): + safe_save(get_state_dict(model), model_name) + print(f"✓ Model saved to: {model_name}") + + if cfg.training.get("log_artifacts", True): + output_dir = cfg.get("output_dir", "outputs") + os.makedirs(output_dir, exist_ok=True) + + plt.figure(figsize=(10, 4)) + plt.plot(losses) + plt.xlabel("Iteration") + plt.ylabel("Loss") + plt.title("Training Loss Over Time") + plt.grid(True) + plt.tight_layout() + plt.savefig(os.path.join(output_dir, "loss_curve.png")) + plt.close() + + if cfg.training.get("generate_samples", True): + num_samples = cfg.training.get("n_samples", 100) + x = T.randn(num_samples, 2) + h_step = cfg.training.step_size + time_grid = T.linspace(0, 1, int(1 / h_step)) + + solver = create_solver(cfg, model, get_preprocess_hook(cfg)) + visualize_moons( + x, + solver=solver, + time_grid=time_grid, + h_step=h_step, + num_plots=cfg.training.get("num_plots", 10), + ) + + +def train_images(cfg: DictConfig, model_name: str): + """Training loop for image datasets (MNIST, Fashion MNIST, CIFAR-10).""" + model = create_model(cfg) + scheduler = create_scheduler(cfg) + path = AffinePath(scheduler=scheduler) + + optim = Adam(get_parameters(model), lr=cfg.optimizer.lr) + lr_scheduler = create_lr_scheduler(cfg, optim) + + dataloader = create_dataloader(cfg) + trainer = create_trainer(cfg, model, dataloader, optim, path, lr_scheduler) + + dataset_name = cfg.dataset.get("type", cfg.dataset.name) + with mlflow.start_run(run_name=dataset_name): + mlflow.log_params(dict(cfg)) + mlflow.log_param("model_name", model_name) + model = trainer.train() + trainer.save_model(model_name) + + if cfg.training.get("log_artifacts", True): + output_dir = cfg.get("output_dir", "outputs") + os.makedirs(output_dir, exist_ok=True) + trainer.plot_loss(output_dir, log_to_mlflow=True) + + if cfg.training.get("generate_samples", True): + solver = create_solver(cfg, model, get_preprocess_hook(cfg)) + trainer.predict(cfg, solver) + + +@hydra.main(version_base=None, config_path="../configs", config_name="config") +def main(cfg: DictConfig): + """Hydra entry point; dispatches on dataset type.""" + print("Configuration:") + print(OmegaConf.to_yaml(cfg)) + + if cfg.get("seed"): + T.manual_seed(cfg.seed) + + model_name = cfg.get("model_name", generate_model_name(cfg)) + print(f"\nModel will be saved as: {model_name}") + + dataset_type = cfg.dataset.get("type", cfg.dataset.name) + if dataset_type == "moons": + train_moons(cfg, model_name) + else: + train_images(cfg, model_name) + + +if __name__ == "__main__": + main() From b51b3431f457519d0ed68485d2ed99842ad5b585 Mon Sep 17 00:00:00 2001 From: Ilia Ilmer Date: Sun, 26 Jul 2026 05:58:10 -0400 Subject: [PATCH 2/4] ignore weights --- .gitignore | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.gitignore b/.gitignore index 55a45af..ba68f73 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,5 @@ outputs *.md *.db !README.md + +weights/ From 37f3628797986e69a54b2633abffd8104e14e44b Mon Sep 17 00:00:00 2001 From: Ilia Ilmer Date: Sun, 26 Jul 2026 06:58:56 -0400 Subject: [PATCH 3/4] deadcode --- scripts/generate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/generate.py b/scripts/generate.py index 0df2ba0..bd7d12b 100644 --- a/scripts/generate.py +++ b/scripts/generate.py @@ -321,7 +321,7 @@ def jit_step(h, t, x): # Add prediction label if available if predictions is not None: - class_idx, class_name, confidence = predictions[idx] + _, class_name, confidence = predictions[idx] label = f"{class_name}\n{confidence:.1%}" ax.set_title(label, fontsize=12, pad=4, fontweight="bold") From 65e074b860699c2e3f37609367c653145fd239eb Mon Sep 17 00:00:00 2001 From: Ilia Ilmer Date: Sun, 26 Jul 2026 07:00:19 -0400 Subject: [PATCH 4/4] modify pyproject --- pyproject.toml | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1ecc612..fb18c28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,12 +40,13 @@ target-version = "py311" [tool.ruff.lint] select = [ - "E", # pycodestyle errors - "W", # pycodestyle warnings - "F", # pyflakes - "I", # isort - "B", # flake8-bugbear - "UP", # pyupgrade + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "B", # flake8-bugbear + "UP", # pyupgrade + "ERA", # commented out code ] ignore = [ "E501", # line too long (handled by formatter)