diff options
Diffstat (limited to 'worldalign/energy_infer.py')
| -rw-r--r-- | worldalign/energy_infer.py | 249 |
1 files changed, 249 insertions, 0 deletions
diff --git a/worldalign/energy_infer.py b/worldalign/energy_infer.py new file mode 100644 index 0000000..d5a67b6 --- /dev/null +++ b/worldalign/energy_infer.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +import torch.nn.functional as F + +from .common import read_json, seed_everything, write_json +from .energy import ( + projection_quantile_target, + prototype_manifold_energy, + relation_field_energy, + retrieval_metrics, + sliced_distribution_energy, + standardized_relation, +) +from .io import load_feature_pair, select_rows + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--manifest", default="artifacts/manifest.json") + parser.add_argument("--vision", default="artifacts/vision.pt") + parser.add_argument("--text", default="artifacts/text.pt") + parser.add_argument( + "--text-orbits", + help=( + "Optional multi-description feature cache. When supplied, each " + "language particle is the normalized mean of its observation " + "orbit." + ), + ) + parser.add_argument("--prototypes", default="artifacts/gw.pt") + parser.add_argument("--split", choices=["val", "test"], default="test") + parser.add_argument("--samples", type=int, default=256) + parser.add_argument("--steps", type=int, default=100) + parser.add_argument("--lr", type=float, default=0.03) + parser.add_argument("--projections", type=int, default=128) + parser.add_argument("--relation-weight-start", type=float, default=0.4) + parser.add_argument("--relation-weight-end", type=float, default=2.0) + parser.add_argument("--conditional-weight-start", type=float, default=0.05) + parser.add_argument("--conditional-weight-end", type=float, default=0.2) + parser.add_argument("--distribution-weight", type=float, default=80.0) + parser.add_argument("--manifold-weight", type=float, default=1.0) + parser.add_argument("--noise", type=float, default=0.01) + parser.add_argument( + "--shuffle-visual-energy", + action="store_true", + help=( + "Evaluation control: permute image particles before constructing " + "the energy while leaving evaluation rows unchanged." + ), + ) + parser.add_argument("--device", default="cuda:1") + parser.add_argument("--seed", type=int, default=20260729) + parser.add_argument( + "--output", default="artifacts/energy_free_test.pt" + ) + parser.add_argument( + "--metrics-output", default="artifacts/energy_free_test.json" + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + seed_everything(args.seed) + manifest = read_json(args.manifest) + vision, text, vision_lookup, text_lookup = load_feature_pair( + args.vision, args.text + ) + rows = manifest[args.split][: args.samples] + visual = select_rows( + vision["features"], vision_lookup, rows + ).to(args.device) + paired_text = select_rows( + text["features"], text_lookup, rows + ).to(args.device) + text_population = select_rows( + text["features"], text_lookup, manifest["text_only_train"] + ).to(args.device) + if args.text_orbits: + orbit_state = torch.load( + args.text_orbits, map_location="cpu", weights_only=False + ) + orbit_lookup = { + int(row): index + for index, row in enumerate(orbit_state["rows"]) + } + orbit_mean = F.normalize( + orbit_state["features"].float().mean(1), dim=-1 + ) + paired_text = select_rows( + orbit_mean, orbit_lookup, rows + ).to(args.device) + text_population = select_rows( + orbit_mean, orbit_lookup, manifest["text_only_train"] + ).to(args.device) + if args.shuffle_visual_energy: + control_generator = torch.Generator( + device=args.device + ).manual_seed(args.seed + 10_000) + visual = visual[ + torch.randperm( + len(visual), + generator=control_generator, + device=args.device, + ) + ] + prototype_state = torch.load( + args.prototypes, map_location="cpu", weights_only=False + ) + prototypes = prototype_state["text_centers"].to(args.device) + if prototypes.shape[-1] != text_population.shape[-1]: + raise ValueError( + "Prototype dimension does not match text features; use the GW " + "cache built from the selected text backbone." + ) + + generator = torch.Generator(device=args.device).manual_seed(args.seed) + initial_indices = torch.randperm( + len(text_population), + generator=generator, + device=args.device, + )[: len(visual)] + initial = text_population[initial_indices].clone() + initial = initial + args.noise * torch.randn( + initial.shape, generator=generator, device=args.device + ) + particles = torch.nn.Parameter(F.normalize(initial, dim=-1)) + directions, target_quantiles = projection_quantile_target( + text_population, + particles=len(particles), + projections=args.projections, + generator=generator, + ) + visual_relation, visual_standardized = standardized_relation(visual) + optimizer = torch.optim.Adam([particles], lr=args.lr) + + def energy_values() -> tuple[torch.Tensor, ...]: + relation, conditional = relation_field_energy( + visual_relation, visual_standardized, particles + ) + distribution = sliced_distribution_energy( + particles, directions, target_quantiles + ) + manifold = prototype_manifold_energy(particles, prototypes) + return relation, conditional, distribution, manifold + + history: list[dict] = [] + initial_cpu = F.normalize(particles.detach(), dim=-1).cpu() + for step in range(args.steps + 1): + relation, conditional, distribution, manifold = energy_values() + progress = min(step / max(args.steps, 1), 1.0) + relation_weight = ( + args.relation_weight_start + + progress + * (args.relation_weight_end - args.relation_weight_start) + ) + conditional_weight = ( + args.conditional_weight_start + + progress + * ( + args.conditional_weight_end + - args.conditional_weight_start + ) + ) + loss = ( + relation_weight * relation + + conditional_weight * conditional + + args.distribution_weight * distribution + + args.manifold_weight * manifold + ) + if step % 20 == 0 or step == args.steps: + history.append( + { + "step": step, + "total": float(loss.detach()), + "relation": float(relation.detach()), + "conditional": float(conditional.detach()), + "distribution": float(distribution.detach()), + "manifold": float(manifold.detach()), + "paired_evaluation_only": retrieval_metrics( + particles.detach(), paired_text + ), + } + ) + print(json.dumps(history[-1])) + if step == args.steps: + break + optimizer.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_([particles], 2.0) + optimizer.step() + with torch.no_grad(): + particles.copy_(F.normalize(particles, dim=-1)) + + with torch.no_grad(): + oracle_relation, oracle_conditional = relation_field_energy( + visual_relation, visual_standardized, paired_text + ) + oracle_distribution = sliced_distribution_energy( + paired_text, directions, target_quantiles + ) + oracle_manifold = prototype_manifold_energy( + paired_text, prototypes + ) + result = { + "protocol": ( + "No cross-modal map and no image-text pair is used by the energy " + "or optimizer. Paired text is loaded only for trajectory and " + "oracle diagnostics; the final step is fixed by CLI arguments." + ), + "mode": "free_language_latent_particles", + "split": args.split, + "rows": rows, + "vision_model": vision["model"], + "text_model": text["model"], + "text_observation": ( + "multi-description orbit mean" + if args.text_orbits + else "single description" + ), + "args": vars(args), + "history": history, + "oracle_energy_components": { + "relation": float(oracle_relation), + "conditional": float(oracle_conditional), + "distribution": float(oracle_distribution), + "manifold": float(oracle_manifold), + }, + } + state = { + **result, + "initial_particles": initial_cpu, + "final_particles": F.normalize( + particles.detach(), dim=-1 + ).cpu(), + } + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + torch.save(state, args.output) + write_json(args.metrics_output, result) + print(f"Wrote {args.output} and {args.metrics_output}") + + +if __name__ == "__main__": + main() |