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()