| from __future__ import annotations |
|
|
| import json |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import trackio |
| from model import ConditionalNeuralProcess, parameter_count |
| from safetensors.torch import save_file |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "neural-process-pocket" |
| DATA_DIR = PROJECT_DIR / "data" |
| SEED = 2141 |
|
|
|
|
| def sample_tasks( |
| tasks: int, |
| context_points: int, |
| target_points: int, |
| seed: int, |
| ) -> tuple[torch.Tensor, ...]: |
| generator = torch.Generator().manual_seed(seed) |
| amplitude = 0.1 + 4.9 * torch.rand(tasks, 1, 1, generator=generator) |
| phase = 2 * math.pi * torch.rand(tasks, 1, 1, generator=generator) |
| frequency = 0.8 + 0.4 * torch.rand(tasks, 1, 1, generator=generator) |
| context_x = -5 + 10 * torch.rand( |
| tasks, context_points, 1, generator=generator |
| ) |
| target_x = -5 + 10 * torch.rand(tasks, target_points, 1, generator=generator) |
| context_y = amplitude * torch.sin(frequency * context_x + phase) |
| target_y = amplitude * torch.sin(frequency * target_x + phase) |
| context_y += 0.03 * torch.randn(context_y.shape, generator=generator) |
| return context_x, context_y, target_x, target_y, amplitude, phase, frequency |
|
|
|
|
| def gaussian_nll( |
| mean: torch.Tensor, |
| standard_deviation: torch.Tensor, |
| target: torch.Tensor, |
| ) -> torch.Tensor: |
| variance = standard_deviation.square() |
| return ( |
| 0.5 * ((target - mean).square() / variance) |
| + standard_deviation.log() |
| + 0.5 * math.log(2 * math.pi) |
| ).mean() |
|
|
|
|
| def rbf_gp( |
| context_x: np.ndarray, |
| context_y: np.ndarray, |
| target_x: np.ndarray, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| def kernel(left: np.ndarray, right: np.ndarray) -> np.ndarray: |
| return np.exp(-0.5 * (left[:, None] - right[None, :]) ** 2) |
|
|
| covariance = kernel(context_x, context_x) + np.eye(len(context_x)) * 0.01 |
| cross = kernel(context_x, target_x) |
| solve_y = np.linalg.solve(covariance, context_y) |
| mean = cross.T @ solve_y |
| solve_cross = np.linalg.solve(covariance, cross) |
| variance = 1.0 - np.sum(cross * solve_cross, axis=0) |
| return mean, np.sqrt(np.clip(variance + 0.03**2, 0.03**2, None)) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model: ConditionalNeuralProcess, tasks: int = 500) -> tuple[dict, list]: |
| context_x, context_y, target_x, target_y, amplitude, phase, frequency = sample_tasks( |
| tasks, 5, 100, SEED + 20_000 |
| ) |
| mean, std = model(context_x, context_y, target_x) |
| error = mean - target_y |
| cnp = { |
| "rmse": float(error.square().mean().sqrt()), |
| "gaussian_nll": float(gaussian_nll(mean, std, target_y)), |
| "coverage_90": float( |
| ((target_y >= mean - 1.645 * std) & (target_y <= mean + 1.645 * std)) |
| .float() |
| .mean() |
| ), |
| } |
| gp_errors = [] |
| gp_nll = [] |
| gp_covered = [] |
| rows = [] |
| for index in range(tasks): |
| cx = context_x[index, :, 0].numpy() |
| cy = context_y[index, :, 0].numpy() |
| tx = target_x[index, :, 0].numpy() |
| ty = target_y[index, :, 0].numpy() |
| gp_mean, gp_std = rbf_gp(cx, cy, tx) |
| gp_errors.extend((gp_mean - ty).tolist()) |
| gp_nll.extend( |
| ( |
| 0.5 * ((ty - gp_mean) / gp_std) ** 2 |
| + np.log(gp_std) |
| + 0.5 * np.log(2 * np.pi) |
| ).tolist() |
| ) |
| gp_covered.extend( |
| ((ty >= gp_mean - 1.645 * gp_std) & (ty <= gp_mean + 1.645 * gp_std)) |
| .astype(float) |
| .tolist() |
| ) |
| rows.append( |
| { |
| "amplitude": float(amplitude[index, 0, 0]), |
| "phase": float(phase[index, 0, 0]), |
| "frequency": float(frequency[index, 0, 0]), |
| "context_x": cx.tolist(), |
| "context_y": cy.tolist(), |
| } |
| ) |
| gp_errors_array = np.asarray(gp_errors) |
| report = { |
| "conditional_neural_process": cnp, |
| "fixed_rbf_gaussian_process": { |
| "rmse": float(np.sqrt(np.mean(gp_errors_array**2))), |
| "gaussian_nll": float(np.mean(gp_nll)), |
| "coverage_90": float(np.mean(gp_covered)), |
| }, |
| "tasks": tasks, |
| "context_points_per_task": 5, |
| "targets_per_task": 100, |
| } |
| return report, rows |
|
|
|
|
| def main() -> None: |
| torch.manual_seed(SEED) |
| torch.set_num_threads(1) |
| model = ConditionalNeuralProcess() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5) |
| best = float("inf") |
| best_state = None |
| best_step = 0 |
| trackio.init( |
| project="neural-process-pocket", |
| name="conditional-neural-process-v1", |
| config={ |
| "parameters": parameter_count(model), |
| "training_steps": 5_000, |
| "context_points": 5, |
| "tasks_per_step": 64, |
| }, |
| ) |
| for step in range(1, 5_001): |
| context_x, context_y, target_x, target_y, *_ = sample_tasks( |
| 64, 5, 50, SEED + step |
| ) |
| mean, std = model(context_x, context_y, target_x) |
| loss = gaussian_nll(mean, std, target_y) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 5) |
| optimizer.step() |
| if step % 250 == 0: |
| validation, _ = evaluate(model, tasks=100) |
| validation_nll = validation["conditional_neural_process"]["gaussian_nll"] |
| trackio.log( |
| { |
| "training_step": step, |
| "training_nll": float(loss.detach()), |
| **{ |
| f"validation_{key}": value |
| for key, value in validation[ |
| "conditional_neural_process" |
| ].items() |
| }, |
| } |
| ) |
| if validation_nll < best: |
| best = validation_nll |
| best_step = step |
| best_state = { |
| name: value.detach().cpu().clone() |
| for name, value in model.state_dict().items() |
| } |
| assert best_state is not None |
| model.load_state_dict(best_state) |
| benchmark, rows = evaluate(model) |
| report = { |
| "model": "Neural Process Pocket", |
| "parameters": parameter_count(model), |
| "best_step": best_step, |
| "benchmark": benchmark, |
| } |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors") |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| pd.DataFrame(rows).to_parquet(DATA_DIR / "heldout_tasks.parquet", index=False) |
| trackio.log( |
| { |
| "test_rmse": benchmark["conditional_neural_process"]["rmse"], |
| "test_nll": benchmark["conditional_neural_process"]["gaussian_nll"], |
| "test_coverage_90": benchmark["conditional_neural_process"][ |
| "coverage_90" |
| ], |
| } |
| ) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
