Use precomputed exposure embeddings

encode_exposure_cache.py

@@ -0,0 +1,147 @@
+"""Encode cached exposure windows once for embedding-only model training."""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from backbones import TimesNetExposureEncoder
+from dataset import ExposureCache
+from train_util import configure_torch_for_training, resolve_device
+
+
+class ExposureEncodingDataset(Dataset):
+    def __init__(self, cache: ExposureCache):
+        self.cache = cache
+
+    def __len__(self) -> int:
+        return len(self.cache.daily)
+
+    def __getitem__(self, index: int) -> tuple[int, torch.Tensor, torch.Tensor]:
+        return (
+            index,
+            torch.from_numpy(
+                np.array(self.cache.daily[index], dtype=np.float32, copy=True)
+            ),
+            torch.from_numpy(
+                np.array(self.cache.monthly[index], dtype=np.float32, copy=True)
+            ),
+        )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Precompute exposure embeddings from an autoencoder checkpoint"
+    )
+    parser.add_argument("--exposure_cache_dir", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--output_file", default=None)
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument(
+        "--output_dtype", choices=["float16", "float32"], default="float16"
+    )
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+    if args.batch_size <= 0:
+        parser.error("--batch_size must be positive")
+    return args
+
+
+def main() -> None:
+    args = parse_args()
+    device = resolve_device(args.device)
+    configure_torch_for_training(device)
+    cache_dir = Path(args.exposure_cache_dir)
+    output_path = (
+        Path(args.output_file)
+        if args.output_file
+        else cache_dir / "exposure_embeddings.npy"
+    )
+    if output_path.exists() and not args.overwrite:
+        raise FileExistsError(f"{output_path} exists; pass --overwrite")
+
+    checkpoint_data = torch.load(args.checkpoint, map_location="cpu")
+    model_cfg = checkpoint_data["model_config"]
+    normalization = checkpoint_data["normalization"]
+    encoder = TimesNetExposureEncoder(
+        n_embd=int(model_cfg["n_embd"]),
+        d_model=model_cfg["d_model"],
+        n_layers=int(model_cfg["n_layers"]),
+        top_k=int(model_cfg["top_k"]),
+        n_backbone_blocks=int(model_cfg["n_backbone_blocks"]),
+        backbone_kernel_size=int(model_cfg["backbone_kernel_size"]),
+        backbone_expansion=float(model_cfg["backbone_expansion"]),
+        dropout=float(model_cfg["dropout"]),
+        use_gate=True,
+    )
+    encoder.load_state_dict(checkpoint_data["encoder_state_dict"], strict=True)
+    encoder.to(device).eval()
+
+    cache = ExposureCache(cache_dir)
+    loader_kwargs = {
+        "batch_size": args.batch_size,
+        "shuffle": False,
+        "num_workers": args.num_workers,
+        "pin_memory": device.type == "cuda",
+        "persistent_workers": args.num_workers > 0,
+    }
+    loader = DataLoader(ExposureEncodingDataset(cache), **loader_kwargs)
+    output_dtype = np.float16 if args.output_dtype == "float16" else np.float32
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output = np.lib.format.open_memmap(
+        output_path,
+        mode="w+",
+        dtype=output_dtype,
+        shape=(len(cache.daily), int(model_cfg["n_embd"])),
+    )
+    stats = {
+        key: torch.as_tensor(normalization[key], device=device).view(1, 1, -1)
+        for key in ("daily_mean", "daily_std", "monthly_mean", "monthly_std")
+    }
+
+    with torch.inference_mode():
+        for indices, daily, monthly in tqdm(loader, desc="Encoding exposure"):
+            daily = daily.to(device, non_blocking=True)
+            monthly = monthly.to(device, non_blocking=True)
+            daily_mask = torch.isfinite(daily)
+            monthly_mask = torch.isfinite(monthly)
+            has_observation = (
+                daily_mask.flatten(1).any(dim=1)
+                | monthly_mask.flatten(1).any(dim=1)
+            )
+            daily = (
+                torch.nan_to_num(daily) - stats["daily_mean"]
+            ) / stats["daily_std"]
+            monthly = (
+                torch.nan_to_num(monthly) - stats["monthly_mean"]
+            ) / stats["monthly_std"]
+            daily = daily * daily_mask
+            monthly = monthly * monthly_mask
+            encoded = encoder(daily, monthly, daily_mask, monthly_mask)
+            encoded = encoded * has_observation.unsqueeze(1)
+            output[indices.numpy()] = encoded.float().cpu().numpy().astype(
+                output_dtype, copy=False
+            )
+    output.flush()
+    metadata = {
+        "checkpoint": str(Path(args.checkpoint).resolve()),
+        "output_file": str(output_path.resolve()),
+        "rows": len(cache.daily),
+        "embedding_dim": int(model_cfg["n_embd"]),
+        "dtype": args.output_dtype,
+    }
+    output_path.with_suffix(".json").write_text(
+        json.dumps(metadata, indent=2), encoding="utf-8"
+    )
+    print(f"Wrote {len(cache.daily):,} embeddings to {output_path}")
+
+
+if __name__ == "__main__":
+    main()