| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
| import time |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
|
|
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from torch.optim import AdamW |
| from torch.utils.data import DataLoader, WeightedRandomSampler |
| from tqdm import tqdm |
|
|
| from .data import SpectraListDataset, collect_mmu_desi, compute_sample_stats, split_indices, valid_pixel_mask |
| from .metrics import LossConfig, masked_huber, redshift_losses, redshift_metrics |
| from .model import fourier_loglam |
| from .plots import plot_reconstruction_batch, plot_redshift_scatter |
|
|
|
|
| @dataclass |
| class RawCollatorConfig: |
| target_length: int = 4096 |
| min_scale: float = 1e-3 |
| random_mask_ratio: float = 0.0 |
| eval_mask_ratio: float = 0.25 |
| mask_mode: str = "pixel" |
| mask_span_min: int = 16 |
| mask_span_max: int = 64 |
| line_region_percentile: float = 90.0 |
| augment_ood: bool = False |
| crop_prob: float = 0.0 |
| bad_window_prob: float = 0.0 |
| throughput_prob: float = 0.0 |
| noise_prob: float = 0.0 |
| resolution_prob: float = 0.0 |
| downsample_prob: float = 0.0 |
| line_dropout_prob: float = 0.0 |
| span_dropout_prob: float = 0.0 |
| grid_jitter_prob: float = 0.0 |
| grid_shift_frac: float = 0.0 |
| grid_scale_frac: float = 0.0 |
| grid_jitter_warmup_steps: int = 0 |
| redshift_shift: float = 0.0 |
|
|
|
|
| class RawSpectraCollator: |
| def __init__(self, cfg: RawCollatorConfig, train: bool = True, seed: int = 17): |
| self.cfg = cfg |
| self.train = train |
| self.seed = seed |
| self.rng = np.random.default_rng(seed) |
| self.batch_count = 0 |
| self.jitter_scale = 1.0 |
|
|
| def __call__(self, samples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: |
| if self.train: |
| self.batch_count += 1 |
| warmup = max(0, int(self.cfg.grid_jitter_warmup_steps)) |
| self.jitter_scale = min(1.0, self.batch_count / float(warmup)) if warmup > 0 else 1.0 |
| items = [self._prepare_sample(s) for s in samples] |
| x = np.stack([item["x"] for item in items], axis=0).astype(np.float32) |
| valid = np.stack([item["valid"] for item in items], axis=0).astype(np.bool_) |
| loglam = np.stack([item["loglam"] for item in items], axis=0).astype(np.float32) |
| target_flux = np.stack([item["target_flux"] for item in items], axis=0).astype(np.float32) |
| loss_mask = np.stack([item["loss_mask"] for item in items], axis=0).astype(np.bool_) |
| line_weight = np.stack([item["line_weight"] for item in items], axis=0).astype(np.float32) |
| line_region = np.stack([item["line_region"] for item in items], axis=0).astype(np.bool_) |
| z = np.asarray([item["z"] for item in items], dtype=np.float32) |
| y = np.asarray([item["y"] for item in items], dtype=np.float32) |
| zwarn = np.asarray([item["zwarn"] for item in items], dtype=np.bool_) |
| return { |
| "x": torch.from_numpy(x), |
| "valid": torch.from_numpy(valid), |
| "loglam": torch.from_numpy(loglam), |
| "target_flux": torch.from_numpy(target_flux), |
| "loss_mask": torch.from_numpy(loss_mask), |
| "line_weight": torch.from_numpy(line_weight), |
| "line_region": torch.from_numpy(line_region), |
| "z": torch.from_numpy(z), |
| "y": torch.from_numpy(y), |
| "zwarn": torch.from_numpy(zwarn), |
| } |
|
|
| def _prepare_sample(self, sample: dict[str, Any]) -> dict[str, Any]: |
| rng = self.rng if self.train else self._eval_rng(sample) |
| flux = np.asarray(sample["flux"], dtype=np.float32).copy() |
| ivar = np.asarray(sample["ivar"], dtype=np.float32).copy() |
| lam = np.asarray(sample["lambda"], dtype=np.float32) |
| lsf = np.asarray(sample["lsf_sigma"], dtype=np.float32) |
| bad = np.asarray(sample["bad_mask"], dtype=np.bool_).copy() |
|
|
| if self.cfg.augment_ood: |
| bad = self._augment_bad_windows(bad, rng) |
| flux = self._augment_flux_calibration(flux, lam, rng) |
| flux = self._augment_resolution(flux, rng) |
| flux, ivar = self._augment_downsample_resample(flux, ivar, lam, rng) |
| flux = self._augment_noise(flux, ivar, rng) |
|
|
| valid = np.isfinite(flux) & np.isfinite(ivar) & np.isfinite(lam) & (ivar > 0) & (~bad) |
| loglam = np.log(np.clip(lam.astype(np.float64), 1.0, None)).astype(np.float32) |
| if valid.sum() < 16: |
| valid = valid_pixel_mask(sample) |
|
|
| grid_lo = float(np.nanmin(loglam)) |
| grid_hi = float(np.nanmax(loglam)) |
| if ( |
| self.train |
| and self.cfg.grid_jitter_prob > 0 |
| and rng.random() < self.cfg.grid_jitter_prob * self.jitter_scale |
| and math.isfinite(grid_lo) |
| and math.isfinite(grid_hi) |
| and grid_hi > grid_lo |
| ): |
| span = grid_hi - grid_lo |
| shift = float(rng.normal(0.0, max(0.0, self.cfg.grid_shift_frac) * self.jitter_scale)) * span |
| scale_max = max(0.0, self.cfg.grid_scale_frac) * self.jitter_scale |
| scale_delta = float(rng.uniform(-scale_max, scale_max)) |
| scaled_span = span * max(0.50, 1.0 + scale_delta) |
| center = 0.5 * (grid_lo + grid_hi) + shift |
| grid_lo = center - 0.5 * scaled_span |
| grid_hi = center + 0.5 * scaled_span |
| grid = np.linspace(grid_lo, grid_hi, self.cfg.target_length, dtype=np.float32) |
| flux_grid = self._interp_valid(loglam, flux, valid, grid, fill=0.0) |
| ivar_grid = self._interp_valid(loglam, ivar, valid, grid, fill=0.0) |
| lsf_grid = self._interp_valid(loglam, lsf, valid, grid, fill=0.0) |
| valid_grid = np.interp(grid, loglam, valid.astype(np.float32), left=0.0, right=0.0) > 0.5 |
|
|
| center = float(np.nanmedian(flux_grid[valid_grid])) if valid_grid.any() else 0.0 |
| dev = np.abs(flux_grid[valid_grid] - center) if valid_grid.any() else np.asarray([1.0], dtype=np.float32) |
| scale = float(np.nanmedian(dev) * 1.4826) |
| if not math.isfinite(scale) or scale < self.cfg.min_scale: |
| scale = max(float(np.nanmedian(np.abs(flux_grid[valid_grid]))) if valid_grid.any() else 1.0, self.cfg.min_scale) |
|
|
| norm_flux = np.arcsinh((flux_grid - center) / scale).astype(np.float32) |
| norm_ivar = np.log1p(np.maximum(ivar_grid * scale * scale, 0.0)).astype(np.float32) |
| norm_ivar = np.clip(norm_ivar / 8.0, 0.0, 4.0) |
| lsf_norm = np.nan_to_num(lsf_grid / 3.0, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32) |
| loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32) |
|
|
| grad = np.gradient(norm_flux, grid).astype(np.float32) |
| good_grad = np.abs(grad[valid_grid]) |
| grad_scale = float(np.percentile(good_grad, 95)) if len(good_grad) else 1.0 |
| if not math.isfinite(grad_scale) or grad_scale <= 0: |
| grad_scale = 1.0 |
| grad = np.clip(grad / grad_scale, -5.0, 5.0).astype(np.float32) |
| abs_grad = np.abs(grad).astype(np.float32) |
|
|
| target_flux = norm_flux.copy() |
| line_weight = self._line_weights(abs_grad, valid_grid) |
| line_region = self._line_region(abs_grad, valid_grid) |
| corrupt = self._sample_input_dropout(abs_grad, valid_grid, rng) |
| if corrupt.any(): |
| norm_flux = norm_flux.copy() |
| grad = grad.copy() |
| abs_grad = abs_grad.copy() |
| norm_flux[corrupt] = 0.0 |
| grad[corrupt] = 0.0 |
| abs_grad[corrupt] = 0.0 |
|
|
| y = math.log1p(float(sample["z"])) |
| if self.train and self.cfg.redshift_shift > 0: |
| delta = float(self.rng.uniform(-self.cfg.redshift_shift, self.cfg.redshift_shift)) |
| y = max(0.0, y + delta) |
| grid = (grid + delta).astype(np.float32) |
| loglam_norm = ((grid - math.log(6000.0)) / 0.45).astype(np.float32) |
|
|
| x = np.stack( |
| [ |
| norm_flux, |
| norm_ivar, |
| valid_grid.astype(np.float32), |
| lsf_norm, |
| loglam_norm, |
| grad, |
| abs_grad, |
| corrupt.astype(np.float32), |
| ], |
| axis=0, |
| ) |
| return { |
| "x": x, |
| "valid": valid_grid, |
| "loglam": grid, |
| "target_flux": target_flux, |
| "loss_mask": corrupt & valid_grid, |
| "line_weight": line_weight, |
| "line_region": line_region, |
| "z": sample["z"], |
| "y": np.float32(y), |
| "zwarn": sample["zwarn"], |
| } |
|
|
| def _eval_rng(self, sample: dict[str, Any]) -> np.random.Generator: |
| object_id = str(sample.get("object_id", "")) |
| lam = np.asarray(sample["lambda"], dtype=np.float32) |
| key = f"{self.seed}|{object_id}|{float(sample['z']):.8g}|{len(lam)}|{float(lam[0]):.4f}|{float(lam[-1]):.4f}" |
| digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest() |
| return np.random.default_rng(int.from_bytes(digest, "little", signed=False)) |
|
|
| def _interp_valid(self, x: np.ndarray, y: np.ndarray, valid: np.ndarray, x_new: np.ndarray, fill: float) -> np.ndarray: |
| good = valid & np.isfinite(x) & np.isfinite(y) |
| if good.sum() < 2: |
| return np.full_like(x_new, fill, dtype=np.float32) |
| return np.interp(x_new, x[good], y[good], left=fill, right=fill).astype(np.float32) |
|
|
| def _augment_bad_windows(self, bad: np.ndarray, rng: np.random.Generator) -> np.ndarray: |
| out = bad.copy() |
| n = len(out) |
| if rng.random() < self.cfg.crop_prob: |
| frac = float(rng.uniform(0.62, 0.96)) |
| width = max(32, int(n * frac)) |
| start = int(rng.integers(0, max(1, n - width))) |
| keep = np.zeros(n, dtype=np.bool_) |
| keep[start : start + width] = True |
| out |= ~keep |
| if rng.random() < self.cfg.bad_window_prob: |
| for _ in range(int(rng.integers(1, 5))): |
| width = int(rng.integers(max(8, n // 240), max(12, n // 45))) |
| start = int(rng.integers(0, max(1, n - width))) |
| out[start : start + width] = True |
| return out |
|
|
| def _augment_flux_calibration(self, flux: np.ndarray, lam: np.ndarray, rng: np.random.Generator) -> np.ndarray: |
| if rng.random() >= self.cfg.throughput_prob: |
| return flux |
| x = np.linspace(-1.0, 1.0, len(flux), dtype=np.float32) |
| coeff = rng.normal(0.0, [0.05, 0.025, 0.015]).astype(np.float32) |
| curve = 1.0 + coeff[0] * x + coeff[1] * (x * x - 0.33) + coeff[2] * np.sin(np.pi * x) |
| return (flux * np.clip(curve, 0.65, 1.35)).astype(np.float32) |
|
|
| def _augment_noise(self, flux: np.ndarray, ivar: np.ndarray, rng: np.random.Generator) -> np.ndarray: |
| if rng.random() >= self.cfg.noise_prob: |
| return flux |
| sigma = np.zeros_like(flux, dtype=np.float32) |
| good = np.isfinite(ivar) & (ivar > 0) |
| sigma[good] = 1.0 / np.sqrt(np.maximum(ivar[good], 1e-8)) |
| scale = float(rng.uniform(0.15, 0.75)) |
| return (flux + rng.normal(0.0, sigma * scale).astype(np.float32)).astype(np.float32) |
|
|
| def _augment_resolution(self, flux: np.ndarray, rng: np.random.Generator) -> np.ndarray: |
| if rng.random() >= self.cfg.resolution_prob: |
| return flux |
| finite = np.isfinite(flux) |
| fill = float(np.nanmedian(flux[finite])) if finite.any() else 0.0 |
| base = np.nan_to_num(flux, nan=fill, posinf=fill, neginf=fill).astype(np.float32) |
| sigma = float(rng.uniform(0.6, 3.0)) |
| radius = max(2, int(math.ceil(4.0 * sigma))) |
| x = np.arange(-radius, radius + 1, dtype=np.float32) |
| kernel = np.exp(-0.5 * (x / sigma) ** 2) |
| kernel = (kernel / kernel.sum()).astype(np.float32) |
| padded = np.pad(base, (radius, radius), mode="edge") |
| return np.convolve(padded, kernel, mode="valid").astype(np.float32) |
|
|
| def _augment_downsample_resample( |
| self, |
| flux: np.ndarray, |
| ivar: np.ndarray, |
| lam: np.ndarray, |
| rng: np.random.Generator, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| if rng.random() >= self.cfg.downsample_prob: |
| return flux, ivar |
| n = len(flux) |
| if n < 32: |
| return flux, ivar |
| factor = int(rng.choice(np.asarray([2, 3, 4, 6, 8], dtype=np.int64))) |
| offset = int(rng.integers(0, factor)) |
| idx = np.arange(offset, n, factor, dtype=np.int64) |
| if len(idx) < 4: |
| return flux, ivar |
| lam_good = np.asarray(lam[idx], dtype=np.float32) |
| flux_good = np.asarray(flux[idx], dtype=np.float32) |
| ivar_good = np.asarray(ivar[idx], dtype=np.float32) |
| good = np.isfinite(lam_good) & np.isfinite(flux_good) & np.isfinite(ivar_good) |
| if np.count_nonzero(good) < 4: |
| return flux, ivar |
| lam_good = lam_good[good] |
| order = np.argsort(lam_good) |
| lam_good = lam_good[order] |
| flux_good = flux_good[good][order] |
| ivar_good = ivar_good[good][order] |
| flux_out = np.interp(lam, lam_good, flux_good, left=flux_good[0], right=flux_good[-1]).astype(np.float32) |
| ivar_out = np.interp(lam, lam_good, ivar_good, left=0.0, right=0.0).astype(np.float32) |
| ivar_out *= float(rng.uniform(0.25, 0.85)) |
| return flux_out, ivar_out |
|
|
| def _sample_input_dropout(self, abs_grad: np.ndarray, valid: np.ndarray, rng: np.random.Generator) -> np.ndarray: |
| corrupt = np.zeros_like(valid, dtype=np.bool_) |
| if valid.sum() < 16: |
| return corrupt |
| n = len(valid) |
| valid_idx = np.where(valid)[0] |
| ratio = self.cfg.random_mask_ratio if self.train else self.cfg.eval_mask_ratio |
| if ratio > 0: |
| n_rand = max(1, int(round(len(valid_idx) * min(float(ratio), 1.0)))) |
| if self.cfg.mask_mode == "pixel": |
| corrupt[rng.choice(valid_idx, size=min(n_rand, len(valid_idx)), replace=False)] = True |
| else: |
| line_bias = self.cfg.mask_mode in {"line_span", "mixed_span"} |
| self._add_spans_to_mask(corrupt, valid, abs_grad, n_rand, rng, line_bias=line_bias) |
| if self.train and rng.random() < self.cfg.span_dropout_prob: |
| for _ in range(int(rng.integers(1, 4))): |
| width = int(rng.integers(max(4, n // 220), max(8, n // 55))) |
| start = int(rng.integers(0, max(1, n - width))) |
| corrupt[start : start + width] |= valid[start : start + width] |
| if self.train and rng.random() < self.cfg.line_dropout_prob: |
| score = abs_grad.copy() |
| score[~valid] = 0.0 |
| if np.count_nonzero(score) > 0: |
| k = max(4, n // 96) |
| peaks = np.argsort(score)[-k:] |
| for j in peaks: |
| width = int(rng.integers(max(2, n // 900), max(4, n // 280))) |
| lo = max(0, int(j) - width) |
| hi = min(n, int(j) + width + 1) |
| corrupt[lo:hi] |= valid[lo:hi] |
| return corrupt & valid |
|
|
| def _add_spans_to_mask( |
| self, |
| corrupt: np.ndarray, |
| valid: np.ndarray, |
| abs_grad: np.ndarray, |
| target_count: int, |
| rng: np.random.Generator, |
| *, |
| line_bias: bool, |
| ) -> None: |
| valid_idx = np.where(valid)[0] |
| if len(valid_idx) == 0: |
| return |
| lo_w = max(1, int(self.cfg.mask_span_min)) |
| hi_w = max(lo_w + 1, int(self.cfg.mask_span_max) + 1) |
| probs = None |
| if line_bias: |
| score = abs_grad[valid_idx].astype(np.float64) |
| positive = score[np.isfinite(score) & (score > 0)] |
| scale = float(np.percentile(positive, 90)) if len(positive) else 1.0 |
| if not math.isfinite(scale) or scale <= 0: |
| scale = 1.0 |
| score = np.clip(score / scale, 0.0, 5.0) + 0.05 |
| probs = score / score.sum() |
| max_tries = max(32, target_count * 4) |
| tries = 0 |
| while int(np.count_nonzero(corrupt & valid)) < target_count and tries < max_tries: |
| tries += 1 |
| center = int(rng.choice(valid_idx, p=probs)) |
| width = int(rng.integers(lo_w, hi_w)) |
| lo = max(0, center - width // 2) |
| hi = min(len(valid), lo + width) |
| corrupt[lo:hi] |= valid[lo:hi] |
|
|
| def _line_weights(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray: |
| weight = np.ones_like(abs_grad, dtype=np.float32) |
| if valid.sum() < 16: |
| return weight |
| scale = float(np.percentile(abs_grad[valid], 90)) |
| if math.isfinite(scale) and scale > 0: |
| weight += 2.0 * np.clip(abs_grad / scale, 0.0, 2.0) |
| weight[~valid] = 1.0 |
| return np.clip(weight, 1.0, 5.0).astype(np.float32) |
|
|
| def _line_region(self, abs_grad: np.ndarray, valid: np.ndarray) -> np.ndarray: |
| region = np.zeros_like(valid, dtype=np.bool_) |
| if valid.sum() < 16: |
| return region |
| pct = min(max(float(self.cfg.line_region_percentile), 0.0), 100.0) |
| thresh = float(np.percentile(abs_grad[valid], pct)) |
| if math.isfinite(thresh) and thresh > 0: |
| region = (abs_grad >= thresh) & valid |
| return region.astype(np.bool_) |
|
|
|
|
| class ConvBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 7, stride: int = 1, dropout: float = 0.0): |
| super().__init__() |
| padding = kernel_size // 2 |
| self.net = nn.Sequential( |
| nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False), |
| nn.BatchNorm1d(out_channels), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False), |
| nn.BatchNorm1d(out_channels), |
| ) |
| self.skip = ( |
| nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) |
| if stride != 1 or in_channels != out_channels |
| else nn.Identity() |
| ) |
| self.act = nn.GELU() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.act(self.net(x) + self.skip(x)) |
|
|
|
|
| class LayerScaleEncoderLayer(nn.Module): |
| def __init__(self, d_model: int, heads: int, dropout: float, layerscale_init: float): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(d_model) |
| self.self_attn = nn.MultiheadAttention(d_model, heads, dropout=dropout, batch_first=True) |
| self.dropout1 = nn.Dropout(dropout) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.linear1 = nn.Linear(d_model, d_model * 4) |
| self.linear2 = nn.Linear(d_model * 4, d_model) |
| self.dropout = nn.Dropout(dropout) |
| self.dropout2 = nn.Dropout(dropout) |
| self.act = nn.GELU() |
| init = float(layerscale_init) |
| self.ls1 = nn.Parameter(torch.full((d_model,), init)) |
| self.ls2 = nn.Parameter(torch.full((d_model,), init)) |
|
|
| def forward( |
| self, |
| src: torch.Tensor, |
| src_mask: torch.Tensor | None = None, |
| src_key_padding_mask: torch.Tensor | None = None, |
| is_causal: bool = False, |
| ) -> torch.Tensor: |
| q = self.norm1(src) |
| attn, _ = self.self_attn( |
| q, |
| q, |
| q, |
| attn_mask=src_mask, |
| key_padding_mask=src_key_padding_mask, |
| need_weights=False, |
| is_causal=is_causal, |
| ) |
| src = src + self.ls1 * self.dropout1(attn) |
| ff = self.linear2(self.dropout(self.act(self.linear1(self.norm2(src))))) |
| return src + self.ls2 * self.dropout2(ff) |
|
|
|
|
| class HybridSpecZ(nn.Module): |
| def __init__( |
| self, |
| in_channels: int = 8, |
| d_model: int = 256, |
| conv_width: int = 128, |
| layers: int = 5, |
| heads: int = 8, |
| dropout: float = 0.1, |
| fourier_freqs: int = 32, |
| z_bins: int = 64, |
| y_min: float = 0.0, |
| y_max: float = math.log1p(6.0), |
| prediction_mode: str = "regression", |
| bin_temperature: float = 1.0, |
| residual_scale: float = 0.06, |
| candidate_topk: int = 5, |
| stem_stride: int = 8, |
| rec_hidden_mult: int = 0, |
| rec_refine_width: int = 16, |
| rec_refine_kernel: int = 5, |
| layerscale_init: float = 0.0, |
| ): |
| super().__init__() |
| allowed_modes = { |
| "regression", |
| "softbin", |
| "hybrid", |
| "bin_residual", |
| "ranked_bin_residual", |
| "candidate_rerank", |
| "calibrated_bin_residual", |
| } |
| if prediction_mode not in allowed_modes: |
| raise ValueError(f"prediction_mode must be one of {sorted(allowed_modes)}, got {prediction_mode!r}") |
| self.fourier_freqs = fourier_freqs |
| self.z_bins = z_bins |
| self.y_min = y_min |
| self.y_max = y_max |
| self.prediction_mode = prediction_mode |
| self.bin_temperature = bin_temperature |
| self.residual_scale = residual_scale |
| self.candidate_topk = max(1, min(int(candidate_topk), z_bins)) |
| if stem_stride not in {4, 8}: |
| raise ValueError(f"stem_stride must be 4 or 8, got {stem_stride}") |
| self.stem_stride = int(stem_stride) |
| self.rec_pixels_per_token = int(stem_stride) |
| self.stride_stages = int(round(math.log2(self.stem_stride))) |
| bin_width = (y_max - y_min) / z_bins |
| centers = torch.linspace(y_min + 0.5 * bin_width, y_max - 0.5 * bin_width, z_bins) |
| self.register_buffer("z_bin_centers", centers, persistent=False) |
|
|
| if self.stem_stride == 8: |
| self.stem = nn.Sequential( |
| ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5), |
| ConvBlock(conv_width, conv_width, kernel_size=7, stride=2, dropout=dropout * 0.5), |
| ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5), |
| ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5), |
| ) |
| else: |
| self.stem = nn.Sequential( |
| ConvBlock(in_channels, conv_width, kernel_size=9, stride=2, dropout=dropout * 0.5), |
| ConvBlock(conv_width, d_model, kernel_size=7, stride=2, dropout=dropout * 0.5), |
| ConvBlock(d_model, d_model, kernel_size=5, stride=1, dropout=dropout * 0.5), |
| ) |
| self.pos_proj = nn.Sequential(nn.Linear(fourier_freqs * 2, d_model), nn.LayerNorm(d_model)) |
| self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) |
| |
| self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) |
|
|
| if layerscale_init > 0: |
| enc_layer = LayerScaleEncoderLayer(d_model, heads, dropout, layerscale_init) |
| else: |
| enc_layer = nn.TransformerEncoderLayer( |
| d_model=d_model, |
| nhead=heads, |
| dim_feedforward=d_model * 4, |
| dropout=dropout, |
| batch_first=True, |
| norm_first=True, |
| activation="gelu", |
| ) |
| self.encoder = nn.TransformerEncoder(enc_layer, num_layers=layers) |
| self.pool_gate = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, 1)) |
| head_dim = d_model * 5 |
| self.z_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, 2)) |
| self.z_bin_head = nn.Sequential(nn.LayerNorm(head_dim), nn.Linear(head_dim, d_model), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_model, z_bins)) |
| self.z_candidate_head = nn.Sequential( |
| nn.LayerNorm(head_dim), |
| nn.Linear(head_dim, d_model), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(d_model, z_bins), |
| ) |
| self.z_rerank_head = nn.Sequential( |
| nn.LayerNorm(head_dim + 3), |
| nn.Linear(head_dim + 3, max(64, d_model // 2)), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(max(64, d_model // 2), 1), |
| ) |
| self.z_calib_head = nn.Sequential( |
| nn.LayerNorm(head_dim + 3), |
| nn.Linear(head_dim + 3, max(64, d_model // 2)), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(max(64, d_model // 2), 1), |
| ) |
| nn.init.zeros_(self.z_calib_head[-1].weight) |
| nn.init.zeros_(self.z_calib_head[-1].bias) |
| if rec_hidden_mult > 0: |
| rec_hidden = int(d_model * rec_hidden_mult) |
| self.rec_head = nn.Sequential( |
| nn.LayerNorm(d_model), |
| nn.Linear(d_model, rec_hidden), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(rec_hidden, self.rec_pixels_per_token), |
| ) |
| else: |
| self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, self.rec_pixels_per_token)) |
| rec_pad = int(rec_refine_kernel) // 2 |
| self.rec_refine = nn.Sequential( |
| nn.Conv1d(1, rec_refine_width, kernel_size=rec_refine_kernel, padding=rec_pad), |
| nn.GELU(), |
| nn.Conv1d(rec_refine_width, 1, kernel_size=rec_refine_kernel, padding=rec_pad), |
| ) |
|
|
| def forward(self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor) -> dict[str, torch.Tensor]: |
| bsz = x.shape[0] |
| h = self.stem(x).transpose(1, 2) |
| tok_valid = valid.float().unsqueeze(1) |
| tok_loglam = loglam.unsqueeze(1) |
| for _ in range(self.stride_stages): |
| tok_valid = F.avg_pool1d(tok_valid, kernel_size=2, stride=2, ceil_mode=True) |
| tok_loglam = F.avg_pool1d(tok_loglam, kernel_size=2, stride=2, ceil_mode=True) |
| tok_valid = tok_valid.squeeze(1) > 0.20 |
| tok_loglam = tok_loglam.squeeze(1) |
| if tok_valid.shape[1] != h.shape[1]: |
| tok_valid = tok_valid[:, : h.shape[1]] |
| tok_loglam = tok_loglam[:, : h.shape[1]] |
| h = h[:, : tok_valid.shape[1]] |
|
|
| h = h + self.pos_proj(fourier_loglam(tok_loglam, self.fourier_freqs)) |
| cls = self.cls.expand(bsz, -1, -1) |
| z_query = self.z_query.expand(bsz, -1, -1) |
| src = torch.cat([cls, z_query, h], dim=1) |
| special_valid = torch.ones((bsz, 2), dtype=torch.bool, device=x.device) |
| src_valid = torch.cat([special_valid, tok_valid], dim=1) |
| padding = ~src_valid |
| memory = self.encoder(src, src_key_padding_mask=padding) |
|
|
| spec = memory[:, 2:] |
| spec_valid = src_valid[:, 2:] |
| spec_mask = spec_valid.unsqueeze(-1) |
| rec = self.rec_head(spec).reshape(bsz, -1) |
| rec = rec + self.rec_refine(rec.unsqueeze(1)).squeeze(1) |
| if rec.shape[1] > x.shape[-1]: |
| rec = rec[:, : x.shape[-1]] |
| elif rec.shape[1] < x.shape[-1]: |
| rec = F.pad(rec, (0, x.shape[-1] - rec.shape[1])) |
| denom = spec_valid.float().sum(dim=1).clamp_min(1.0).unsqueeze(-1) |
| mean_pool = (spec * spec_mask.float()).sum(dim=1) / denom |
| max_pool = spec.masked_fill(~spec_mask, -1e4).max(dim=1).values |
| gate_logits = self.pool_gate(spec).squeeze(-1).masked_fill(~spec_valid, -1e4) |
| gate = torch.softmax(gate_logits, dim=1) |
| attn_pool = torch.einsum("bn,bnd->bd", gate, spec) |
| feat = torch.cat([memory[:, 0], memory[:, 1], mean_pool, max_pool, attn_pool], dim=-1) |
| z_params = self.z_head(feat) |
| z_bin_logits = self.z_bin_head(feat) |
| candidate_residual = self.residual_scale * torch.tanh(self.z_candidate_head(feat)) |
| centers = self.z_bin_centers.to(dtype=z_bin_logits.dtype, device=z_bin_logits.device) |
| candidate_y = (centers.unsqueeze(0) + candidate_residual).clamp(self.y_min, self.y_max) |
| topk_logits, topk_bins = torch.topk(z_bin_logits, k=self.candidate_topk, dim=-1) |
| candidate_topk_y = candidate_y.gather(1, topk_bins) |
| rank = torch.linspace(0.0, 1.0, self.candidate_topk, device=x.device, dtype=feat.dtype).view(1, self.candidate_topk, 1) |
| rerank_feat = feat.unsqueeze(1).expand(-1, self.candidate_topk, -1) |
| rerank_in = torch.cat( |
| [ |
| rerank_feat, |
| candidate_topk_y.to(dtype=feat.dtype).unsqueeze(-1), |
| topk_logits.to(dtype=feat.dtype).unsqueeze(-1), |
| rank.expand(bsz, -1, -1), |
| ], |
| dim=-1, |
| ) |
| rerank_logits = self.z_rerank_head(rerank_in).squeeze(-1) |
| rerank_idx = rerank_logits.argmax(dim=-1, keepdim=True) |
| y_reranked = candidate_topk_y.gather(1, rerank_idx).squeeze(1) |
| y_reg = z_params[:, 0] |
| bin_prob = torch.softmax(z_bin_logits / max(self.bin_temperature, 1e-4), dim=-1) |
| y_bin = (bin_prob * self.z_bin_centers.to(dtype=bin_prob.dtype, device=bin_prob.device)).sum(dim=-1) |
| y_ranked = (bin_prob * candidate_y.to(dtype=bin_prob.dtype)).sum(dim=-1) |
| y_legacy_bin_residual = y_bin + self.residual_scale * torch.tanh(y_reg) |
| calib_in = torch.cat( |
| [ |
| feat, |
| y_legacy_bin_residual.to(dtype=feat.dtype).unsqueeze(-1), |
| y_ranked.to(dtype=feat.dtype).unsqueeze(-1), |
| candidate_topk_y[:, 0].to(dtype=feat.dtype).unsqueeze(-1), |
| ], |
| dim=-1, |
| ) |
| y_calibrated = y_legacy_bin_residual + self.residual_scale * torch.tanh(self.z_calib_head(calib_in).squeeze(-1)) |
| if self.prediction_mode == "regression": |
| y_pred = y_reg |
| elif self.prediction_mode == "softbin": |
| y_pred = y_bin |
| elif self.prediction_mode == "hybrid": |
| y_pred = 0.35 * y_reg + 0.65 * y_bin |
| elif self.prediction_mode == "ranked_bin_residual": |
| y_pred = 0.5 * y_legacy_bin_residual + 0.5 * y_ranked |
| elif self.prediction_mode == "candidate_rerank": |
| y_pred = y_reranked |
| elif self.prediction_mode == "calibrated_bin_residual": |
| y_pred = y_calibrated |
| else: |
| y_pred = y_legacy_bin_residual |
| y_pred = y_pred.clamp(self.y_min, self.y_max) |
| return { |
| "rec": rec, |
| "y_mu": y_pred, |
| "y_pred": y_pred, |
| "y_reg": y_reg, |
| "y_bin": y_bin, |
| "y_ranked": y_ranked, |
| "y_top1_candidate": candidate_topk_y[:, 0], |
| "y_reranked": y_reranked, |
| "y_calibrated": y_calibrated, |
| "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0), |
| "z_bin_logits": z_bin_logits, |
| "z_feat": feat, |
| "candidate_y": candidate_y, |
| "candidate_topk_y": candidate_topk_y, |
| "candidate_topk_bins": topk_bins, |
| "candidate_topk_logits": topk_logits, |
| "rerank_logits": rerank_logits, |
| } |
|
|
| def y_to_bin(self, y: torch.Tensor) -> torch.Tensor: |
| scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6) |
| return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1) |
|
|
|
|
| class WavelengthTokenSpecZ(HybridSpecZ): |
| """Transformer encoder over wavelength-conditioned tokens instead of a conv pixel stem.""" |
|
|
| def __init__( |
| self, |
| in_channels: int = 8, |
| d_model: int = 256, |
| conv_width: int = 128, |
| layers: int = 5, |
| heads: int = 8, |
| dropout: float = 0.1, |
| fourier_freqs: int = 32, |
| z_bins: int = 64, |
| y_min: float = 0.0, |
| y_max: float = math.log1p(6.0), |
| prediction_mode: str = "regression", |
| bin_temperature: float = 1.0, |
| residual_scale: float = 0.06, |
| candidate_topk: int = 5, |
| token_stride: int = 8, |
| rec_hidden_mult: int = 0, |
| rec_refine_width: int = 16, |
| rec_refine_kernel: int = 5, |
| layerscale_init: float = 0.0, |
| ): |
| super().__init__( |
| in_channels=in_channels, |
| d_model=d_model, |
| conv_width=conv_width, |
| layers=layers, |
| heads=heads, |
| dropout=dropout, |
| fourier_freqs=fourier_freqs, |
| z_bins=z_bins, |
| y_min=y_min, |
| y_max=y_max, |
| prediction_mode=prediction_mode, |
| bin_temperature=bin_temperature, |
| residual_scale=residual_scale, |
| candidate_topk=candidate_topk, |
| stem_stride=8, |
| rec_hidden_mult=rec_hidden_mult, |
| rec_refine_width=rec_refine_width, |
| rec_refine_kernel=rec_refine_kernel, |
| layerscale_init=layerscale_init, |
| ) |
| self.token_stride = max(1, int(token_stride)) |
| self.rec_pixels_per_token = self.token_stride |
| self.stem = nn.Identity() |
| self.input_proj = nn.Sequential( |
| nn.Linear(in_channels + fourier_freqs * 2, d_model), |
| nn.LayerNorm(d_model), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(d_model, d_model), |
| ) |
| if rec_hidden_mult > 0: |
| rec_hidden = int(d_model * rec_hidden_mult) |
| self.rec_head = nn.Sequential( |
| nn.LayerNorm(d_model), |
| nn.Linear(d_model, rec_hidden), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(rec_hidden, self.rec_pixels_per_token), |
| ) |
| else: |
| self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, self.rec_pixels_per_token)) |
|
|
| def _pool_wavelength_tokens( |
| self, |
| x: torch.Tensor, |
| valid: torch.Tensor, |
| loglam: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| bsz, channels, length = x.shape |
| stride = self.token_stride |
| pad = (-length) % stride |
| if pad: |
| x = F.pad(x, (0, pad)) |
| valid = F.pad(valid.float(), (0, pad)).bool() |
| loglam = torch.cat([loglam, loglam[:, -1:].expand(-1, pad)], dim=1) |
| token_count = x.shape[-1] // stride |
| x_group = x.reshape(bsz, channels, token_count, stride) |
| valid_group = valid.reshape(bsz, 1, token_count, stride).float() |
| loglam_group = loglam.reshape(bsz, 1, token_count, stride) |
| counts = valid_group.sum(dim=-1) |
| denom = counts.clamp_min(1.0) |
| token_x = (x_group * valid_group).sum(dim=-1) / denom |
| token_loglam = (loglam_group * valid_group).sum(dim=-1) / denom |
| fallback_loglam = loglam_group.mean(dim=-1) |
| token_loglam = torch.where(counts > 0, token_loglam, fallback_loglam).squeeze(1) |
| token_valid = counts.squeeze(1) > 0 |
| return token_x.transpose(1, 2), token_valid, token_loglam |
|
|
| def forward(self, x: torch.Tensor, valid: torch.Tensor, loglam: torch.Tensor) -> dict[str, torch.Tensor]: |
| bsz = x.shape[0] |
| token_x, tok_valid, tok_loglam = self._pool_wavelength_tokens(x, valid, loglam) |
| h = self.input_proj(torch.cat([token_x, fourier_loglam(tok_loglam, self.fourier_freqs)], dim=-1)) |
| cls = self.cls.expand(bsz, -1, -1) |
| z_query = self.z_query.expand(bsz, -1, -1) |
| src = torch.cat([cls, z_query, h], dim=1) |
| special_valid = torch.ones((bsz, 2), dtype=torch.bool, device=x.device) |
| src_valid = torch.cat([special_valid, tok_valid], dim=1) |
| padding = ~src_valid |
| memory = self.encoder(src, src_key_padding_mask=padding) |
|
|
| spec = memory[:, 2:] |
| spec_valid = src_valid[:, 2:] |
| spec_mask = spec_valid.unsqueeze(-1) |
| rec = self.rec_head(spec).reshape(bsz, -1) |
| rec = rec + self.rec_refine(rec.unsqueeze(1)).squeeze(1) |
| if rec.shape[1] > x.shape[-1]: |
| rec = rec[:, : x.shape[-1]] |
| elif rec.shape[1] < x.shape[-1]: |
| rec = F.pad(rec, (0, x.shape[-1] - rec.shape[1])) |
| denom = spec_valid.float().sum(dim=1).clamp_min(1.0).unsqueeze(-1) |
| mean_pool = (spec * spec_mask.float()).sum(dim=1) / denom |
| max_pool = spec.masked_fill(~spec_mask, -1e4).max(dim=1).values |
| gate_logits = self.pool_gate(spec).squeeze(-1).masked_fill(~spec_valid, -1e4) |
| gate = torch.softmax(gate_logits, dim=1) |
| attn_pool = torch.einsum("bn,bnd->bd", gate, spec) |
| feat = torch.cat([memory[:, 0], memory[:, 1], mean_pool, max_pool, attn_pool], dim=-1) |
| z_params = self.z_head(feat) |
| z_bin_logits = self.z_bin_head(feat) |
| candidate_residual = self.residual_scale * torch.tanh(self.z_candidate_head(feat)) |
| centers = self.z_bin_centers.to(dtype=z_bin_logits.dtype, device=z_bin_logits.device) |
| candidate_y = (centers.unsqueeze(0) + candidate_residual).clamp(self.y_min, self.y_max) |
| topk_logits, topk_bins = torch.topk(z_bin_logits, k=self.candidate_topk, dim=-1) |
| candidate_topk_y = candidate_y.gather(1, topk_bins) |
| rank = torch.linspace(0.0, 1.0, self.candidate_topk, device=x.device, dtype=feat.dtype).view(1, self.candidate_topk, 1) |
| rerank_feat = feat.unsqueeze(1).expand(-1, self.candidate_topk, -1) |
| rerank_in = torch.cat( |
| [ |
| rerank_feat, |
| candidate_topk_y.to(dtype=feat.dtype).unsqueeze(-1), |
| topk_logits.to(dtype=feat.dtype).unsqueeze(-1), |
| rank.expand(bsz, -1, -1), |
| ], |
| dim=-1, |
| ) |
| rerank_logits = self.z_rerank_head(rerank_in).squeeze(-1) |
| rerank_idx = rerank_logits.argmax(dim=-1, keepdim=True) |
| y_reranked = candidate_topk_y.gather(1, rerank_idx).squeeze(1) |
| y_reg = z_params[:, 0] |
| bin_prob = torch.softmax(z_bin_logits / max(self.bin_temperature, 1e-4), dim=-1) |
| y_bin = (bin_prob * self.z_bin_centers.to(dtype=bin_prob.dtype, device=bin_prob.device)).sum(dim=-1) |
| y_ranked = (bin_prob * candidate_y.to(dtype=bin_prob.dtype)).sum(dim=-1) |
| y_legacy_bin_residual = y_bin + self.residual_scale * torch.tanh(y_reg) |
| calib_in = torch.cat( |
| [ |
| feat, |
| y_legacy_bin_residual.to(dtype=feat.dtype).unsqueeze(-1), |
| y_ranked.to(dtype=feat.dtype).unsqueeze(-1), |
| candidate_topk_y[:, 0].to(dtype=feat.dtype).unsqueeze(-1), |
| ], |
| dim=-1, |
| ) |
| y_calibrated = y_legacy_bin_residual + self.residual_scale * torch.tanh(self.z_calib_head(calib_in).squeeze(-1)) |
| if self.prediction_mode == "regression": |
| y_pred = y_reg |
| elif self.prediction_mode == "softbin": |
| y_pred = y_bin |
| elif self.prediction_mode == "hybrid": |
| y_pred = 0.35 * y_reg + 0.65 * y_bin |
| elif self.prediction_mode == "ranked_bin_residual": |
| y_pred = 0.5 * y_legacy_bin_residual + 0.5 * y_ranked |
| elif self.prediction_mode == "candidate_rerank": |
| y_pred = y_reranked |
| elif self.prediction_mode == "calibrated_bin_residual": |
| y_pred = y_calibrated |
| else: |
| y_pred = y_legacy_bin_residual |
| y_pred = y_pred.clamp(self.y_min, self.y_max) |
| return { |
| "rec": rec, |
| "y_mu": y_pred, |
| "y_pred": y_pred, |
| "y_reg": y_reg, |
| "y_bin": y_bin, |
| "y_ranked": y_ranked, |
| "y_top1_candidate": candidate_topk_y[:, 0], |
| "y_reranked": y_reranked, |
| "y_calibrated": y_calibrated, |
| "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0), |
| "z_bin_logits": z_bin_logits, |
| "z_feat": feat, |
| "candidate_y": candidate_y, |
| "candidate_topk_y": candidate_topk_y, |
| "candidate_topk_bins": topk_bins, |
| "candidate_topk_logits": topk_logits, |
| "rerank_logits": rerank_logits, |
| } |
|
|
|
|
| def move_to_device(batch: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]: |
| return {k: v.to(device, non_blocking=True) if torch.is_tensor(v) else v for k, v in batch.items()} |
|
|
|
|
| def limit_batch_examples(batch: dict[str, torch.Tensor], max_examples: int | None, seen_examples: int) -> dict[str, torch.Tensor] | None: |
| if max_examples is None or max_examples <= 0: |
| return batch |
| remaining = int(max_examples) - int(seen_examples) |
| if remaining <= 0: |
| return None |
| bsz = int(batch["y"].shape[0]) |
| if remaining >= bsz: |
| return batch |
| return {k: v[:remaining] if torch.is_tensor(v) and v.shape[:1] == (bsz,) else v for k, v in batch.items()} |
|
|
|
|
| def load_checkpoint_into_model(model: nn.Module, state: dict[str, torch.Tensor], allow_mismatched: bool = False) -> None: |
| if not allow_mismatched: |
| try: |
| model.load_state_dict(state, strict=True) |
| except RuntimeError: |
| missing, unexpected = model.load_state_dict(state, strict=False) |
| print(f"RESUME_NONSTRICT missing={list(missing)} unexpected={list(unexpected)}") |
| return |
|
|
| target_state = model.state_dict() |
| compatible = {} |
| skipped = [] |
| for key, value in state.items(): |
| target = target_state.get(key) |
| if target is not None and tuple(target.shape) == tuple(value.shape): |
| compatible[key] = value |
| else: |
| skipped.append(key) |
| missing, unexpected = model.load_state_dict(compatible, strict=False) |
| print( |
| "RESUME_FILTERED " |
| f"loaded={len(compatible)} skipped={len(skipped)} " |
| f"missing={list(missing)} unexpected={list(unexpected)} skipped_keys={skipped[:20]}" |
| ) |
|
|
|
|
| def configure_trainable_parameters(model: nn.Module, freeze_mode: str, train_top_layers: int, train_layernorms: bool) -> int: |
| if freeze_mode == "none": |
| for param in model.parameters(): |
| param.requires_grad = True |
| elif freeze_mode == "rerank": |
| for param in model.parameters(): |
| param.requires_grad = False |
| for name, param in model.named_parameters(): |
| if name.startswith("z_rerank_head"): |
| param.requires_grad = True |
| elif freeze_mode == "calib": |
| for param in model.parameters(): |
| param.requires_grad = False |
| for name, param in model.named_parameters(): |
| if name.startswith("z_calib_head"): |
| param.requires_grad = True |
| elif freeze_mode == "adapter": |
| for param in model.parameters(): |
| param.requires_grad = False |
| train_prefixes = ( |
| "stem", |
| "input_proj", |
| "pos_proj", |
| "pool_gate", |
| "z_head", |
| "z_bin_head", |
| "z_candidate_head", |
| "z_rerank_head", |
| "z_calib_head", |
| "rec_head", |
| "rec_refine", |
| "cls", |
| "z_query", |
| ) |
| for name, param in model.named_parameters(): |
| if name.startswith(train_prefixes): |
| param.requires_grad = True |
| if train_layernorms and (".norm" in name or name.endswith("norm.weight") or name.endswith("norm.bias")): |
| param.requires_grad = True |
| layers = getattr(getattr(model, "encoder", None), "layers", None) |
| if layers is not None and train_top_layers > 0: |
| for layer in list(layers)[-int(train_top_layers) :]: |
| for param in layer.parameters(): |
| param.requires_grad = True |
| else: |
| raise ValueError(f"Unknown freeze mode {freeze_mode!r}") |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
|
|
| def replay_loss( |
| student_out: dict[str, torch.Tensor], |
| teacher_out: dict[str, torch.Tensor], |
| batch: dict[str, torch.Tensor], |
| *, |
| y_weight: float, |
| bin_weight: float, |
| clean_only: bool, |
| ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| y_student = student_out.get("y_pred", student_out["y_mu"]).float() |
| y_teacher = teacher_out.get("y_pred", teacher_out["y_mu"]).float().detach() |
| mask = torch.isfinite(batch["y"]) |
| if clean_only: |
| mask = mask & (~batch["zwarn"].bool()) |
| if mask.sum() == 0: |
| zero = y_student.sum() * 0.0 |
| return zero, {"replay_y": zero.detach(), "replay_bin": zero.detach()} |
| replay_y = F.smooth_l1_loss(y_student[mask], y_teacher[mask], beta=0.01) |
| replay_bin = y_student.sum() * 0.0 |
| if bin_weight > 0 and "z_bin_logits" in student_out and "z_bin_logits" in teacher_out: |
| student_logp = F.log_softmax(student_out["z_bin_logits"][mask].float(), dim=-1) |
| teacher_p = F.softmax(teacher_out["z_bin_logits"][mask].float().detach(), dim=-1) |
| replay_bin = F.kl_div(student_logp, teacher_p, reduction="batchmean") |
| total = y_weight * replay_y + bin_weight * replay_bin |
| return total, {"replay_y": replay_y.detach(), "replay_bin": replay_bin.detach()} |
|
|
|
|
| def redshift_total_loss(model: HybridSpecZ, out: dict[str, torch.Tensor], batch: dict[str, torch.Tensor], cfg: LossConfig) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: |
| parts = redshift_losses(model, out, batch["y"], batch["zwarn"], cfg) |
| if "rec" in out and "target_flux" in batch and "loss_mask" in batch: |
| line_weight = batch.get("line_weight") |
| if line_weight is not None: |
| line_weight = line_weight.pow(cfg.line_weight_power) |
| rec = masked_huber(out["rec"], batch["target_flux"], batch["loss_mask"], weight=line_weight) |
| else: |
| rec = parts["z_huber"].sum() * 0.0 |
| total = ( |
| cfg.rec_weight * rec |
| + cfg.z_weight * parts["z_huber"] |
| + cfg.z_bin_weight * parts["z_bin"] |
| + cfg.z_candidate_weight * parts["z_candidate"] |
| + cfg.z_rerank_weight * parts["z_rerank"] |
| + cfg.z_nll_weight * parts["z_nll"] |
| ) |
| metrics = {"loss": total.detach(), "rec": rec.detach(), **{k: v.detach() for k, v in parts.items()}} |
| return total, metrics |
|
|
|
|
| def plot_spectra_batch(path: str | Path, batch: dict[str, torch.Tensor], y_pred: np.ndarray, max_items: int = 4) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| x = batch["x"].detach().cpu().numpy() |
| loglam = batch["loglam"].detach().cpu().numpy() |
| valid = batch["valid"].detach().cpu().numpy() |
| z = batch["z"].detach().cpu().numpy() |
| bsz = min(max_items, x.shape[0]) |
| fig, axes = plt.subplots(bsz, 1, figsize=(13, 3.0 * bsz), squeeze=False) |
| for i in range(bsz): |
| ax = axes[i, 0] |
| wave = np.exp(loglam[i]) |
| good = valid[i].astype(bool) |
| ax.plot(wave[good], x[i, 0, good], color="black", linewidth=0.8, label="input flux") |
| ax.plot(wave[good], x[i, 6, good], color="#1f77b4", linewidth=0.6, alpha=0.55, label="line score") |
| masked = x[i, 7] > 0 |
| if masked.any(): |
| ax.scatter(wave[masked], np.zeros(masked.sum()), s=5, color="#d62728", alpha=0.55, label="redshift dropout") |
| ax.set_title(f"z true={z[i]:.5f} z pred={np.expm1(y_pred[i]):.5f}") |
| ax.set_ylabel("normalized") |
| ax.grid(alpha=0.2) |
| if i == 0: |
| ax.legend(loc="best", fontsize=8) |
| axes[-1, 0].set_xlabel("wavelength Angstrom") |
| fig.tight_layout() |
| fig.savefig(path, dpi=150) |
| plt.close(fig) |
|
|
|
|
| def add_redshift_slice_metrics(metrics: dict[str, float], prefix: str, y_true: np.ndarray, y_pred: np.ndarray) -> None: |
| z_true = np.expm1(y_true) |
| z_pred = np.expm1(y_pred) |
| slices = { |
| "z_lt_0p4": z_true < 0.4, |
| "z_0p4_1p0": (z_true >= 0.4) & (z_true < 1.0), |
| "z_1p0_2p0": (z_true >= 1.0) & (z_true < 2.0), |
| "z_gte_2p0": z_true >= 2.0, |
| } |
| for name, mask in slices.items(): |
| count = int(np.count_nonzero(mask)) |
| metrics[f"{prefix}/{name}_count"] = float(count) |
| if count >= 5: |
| err = z_pred[mask] - z_true[mask] |
| denom = 1.0 + z_true[mask] |
| metrics[f"{prefix}/{name}_mae_z"] = float(np.mean(np.abs(err))) |
| metrics[f"{prefix}/{name}_bias_z"] = float(np.mean(err)) |
| metrics[f"{prefix}/{name}_cat_0p05"] = float(np.mean(np.abs(err / denom) > 0.05)) |
|
|
|
|
| def add_candidate_metrics( |
| metrics: dict[str, float], |
| prefix: str, |
| y_true: np.ndarray, |
| candidate_y: np.ndarray, |
| candidate_bins: np.ndarray | None, |
| *, |
| z_bins: int, |
| y_min: float, |
| y_max: float, |
| ) -> None: |
| if candidate_y.size == 0: |
| return |
| z_true = np.expm1(y_true) |
| z_candidate = np.expm1(candidate_y) |
| abs_dz = np.abs(z_candidate - z_true[:, None]) |
| norm_dz = abs_dz / (1.0 + z_true[:, None]) |
| top_limits = [1, 3, 5] |
| for k in top_limits: |
| kk = min(k, candidate_y.shape[1]) |
| best_abs = np.min(abs_dz[:, :kk], axis=1) |
| best_norm = np.min(norm_dz[:, :kk], axis=1) |
| metrics[f"{prefix}/candidate_top{kk}_best_mae_z"] = float(np.mean(best_abs)) |
| metrics[f"{prefix}/candidate_top{kk}_hit_0p003"] = float(np.mean(best_norm <= 0.003)) |
| metrics[f"{prefix}/candidate_top{kk}_hit_0p01"] = float(np.mean(best_norm <= 0.01)) |
| metrics[f"{prefix}/candidate_top{kk}_hit_0p05"] = float(np.mean(best_norm <= 0.05)) |
| if candidate_bins is not None and candidate_bins.size: |
| scaled = (y_true - y_min) / max(y_max - y_min, 1e-6) |
| true_bins = np.clip((scaled * z_bins).astype(np.int64), 0, z_bins - 1) |
| for k in top_limits: |
| kk = min(k, candidate_bins.shape[1]) |
| metrics[f"{prefix}/candidate_top{kk}_bin_hit"] = float(np.mean(np.any(candidate_bins[:, :kk] == true_bins[:, None], axis=1))) |
|
|
|
|
| @torch.no_grad() |
| def evaluate( |
| model: HybridSpecZ, |
| loader: DataLoader, |
| loss_cfg: LossConfig, |
| device: torch.device, |
| run_dir: Path, |
| step: int, |
| prefix: str = "val", |
| max_batches: int | None = 50, |
| max_examples: int | None = None, |
| ) -> dict[str, float]: |
| model.eval() |
| losses: dict[str, list[float]] = {} |
| y_true_all: list[np.ndarray] = [] |
| y_pred_all: list[np.ndarray] = [] |
| candidate_y_all: list[np.ndarray] = [] |
| candidate_bins_all: list[np.ndarray] = [] |
| y_true_clean: list[np.ndarray] = [] |
| y_pred_clean: list[np.ndarray] = [] |
| candidate_y_clean: list[np.ndarray] = [] |
| candidate_bins_clean: list[np.ndarray] = [] |
| zwarn_all: list[np.ndarray] = [] |
| first_batch = None |
| first_pred = None |
| first_rec = None |
| seen_examples = 0 |
| for bi, batch in enumerate(loader): |
| if max_batches is not None and max_batches > 0 and bi >= max_batches: |
| break |
| batch = limit_batch_examples(batch, max_examples, seen_examples) |
| if batch is None: |
| break |
| seen_examples += int(batch["y"].shape[0]) |
| batch = move_to_device(batch, device) |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| out = model(batch["x"], batch["valid"], batch["loglam"]) |
| _, parts = redshift_total_loss(model, out, batch, loss_cfg) |
| y_pred = out.get("y_pred", out["y_mu"]) |
| for k, v in parts.items(): |
| losses.setdefault(k, []).append(float(v.detach().cpu())) |
| if "rec" in out and "target_flux" in batch and "loss_mask" in batch: |
| rec_err = F.smooth_l1_loss(out["rec"].float(), batch["target_flux"].float(), reduction="none", beta=0.5) |
| loss_mask = batch["loss_mask"].bool() |
| line_region = batch.get("line_region") |
| if line_region is not None: |
| line_mask = loss_mask & line_region.bool() |
| cont_mask = loss_mask & (~line_region.bool()) |
| for name, mask in (("rec_line", line_mask), ("rec_continuum", cont_mask)): |
| denom = mask.float().sum() |
| if float(denom.detach().cpu()) > 0: |
| losses.setdefault(name, []).append(float(((rec_err * mask.float()).sum() / denom.clamp_min(1.0)).detach().cpu())) |
| context_mask = batch["valid"].bool() & (~loss_mask) |
| denom = context_mask.float().sum(dim=1).clamp_min(1.0) |
| baseline = (batch["target_flux"].float() * context_mask.float()).sum(dim=1, keepdim=True) / denom.unsqueeze(1) |
| baseline_err = F.smooth_l1_loss(baseline.expand_as(batch["target_flux"]).float(), batch["target_flux"].float(), reduction="none", beta=0.5) |
| mask_denom = loss_mask.float().sum().clamp_min(1.0) |
| losses.setdefault("rec_mean_baseline", []).append(float(((baseline_err * loss_mask.float()).sum() / mask_denom).detach().cpu())) |
| finite = torch.isfinite(batch["y"]).detach().cpu().numpy() |
| clean = ((~batch["zwarn"].bool()) & torch.isfinite(batch["y"])).detach().cpu().numpy() |
| zw = batch["zwarn"].detach().cpu().numpy().astype(bool) |
| if finite.any(): |
| y_true_all.append(batch["y"].detach().cpu().numpy()[finite]) |
| y_pred_all.append(y_pred.float().detach().cpu().numpy()[finite]) |
| zwarn_all.append(zw[finite]) |
| if "candidate_topk_y" in out: |
| candidate_y_all.append(out["candidate_topk_y"].float().detach().cpu().numpy()[finite]) |
| if "candidate_topk_bins" in out: |
| candidate_bins_all.append(out["candidate_topk_bins"].detach().cpu().numpy()[finite]) |
| if clean.any(): |
| y_true_clean.append(batch["y"].detach().cpu().numpy()[clean]) |
| y_pred_clean.append(y_pred.float().detach().cpu().numpy()[clean]) |
| if "candidate_topk_y" in out: |
| candidate_y_clean.append(out["candidate_topk_y"].float().detach().cpu().numpy()[clean]) |
| if "candidate_topk_bins" in out: |
| candidate_bins_clean.append(out["candidate_topk_bins"].detach().cpu().numpy()[clean]) |
| if first_batch is None: |
| first_batch = {k: v.detach().cpu() if torch.is_tensor(v) else v for k, v in batch.items()} |
| first_pred = y_pred.float().detach().cpu().numpy() |
| if "rec" in out: |
| first_rec = out["rec"].float().detach().cpu().numpy() |
|
|
| metrics = {f"{prefix}/{k}": float(np.mean(v)) for k, v in losses.items()} |
| if y_true_all: |
| y_true = np.concatenate(y_true_all) |
| y_pred = np.concatenate(y_pred_all) |
| for k, v in redshift_metrics(y_true, y_pred).items(): |
| metrics[f"{prefix}/{k}"] = v |
| add_redshift_slice_metrics(metrics, prefix, y_true, y_pred) |
| if candidate_y_all: |
| candidate_y_np = np.concatenate(candidate_y_all) |
| candidate_bins_np = np.concatenate(candidate_bins_all) if candidate_bins_all else None |
| add_candidate_metrics( |
| metrics, |
| prefix, |
| y_true, |
| candidate_y_np, |
| candidate_bins_np, |
| z_bins=model.z_bins, |
| y_min=model.y_min, |
| y_max=model.y_max, |
| ) |
| metrics[f"{prefix}/z_count"] = float(len(y_true)) |
| metrics[f"{prefix}/zwarn_fraction"] = float(np.mean(np.concatenate(zwarn_all))) if zwarn_all else 0.0 |
| plot_redshift_scatter(run_dir / "plots" / f"{prefix}_redshift_step_{step:06d}.png", y_true, y_pred) |
| if y_true_clean: |
| clean_true = np.concatenate(y_true_clean) |
| clean_pred = np.concatenate(y_pred_clean) |
| if len(clean_true) >= 5: |
| for k, v in redshift_metrics(clean_true, clean_pred).items(): |
| metrics[f"{prefix}_clean/{k}"] = v |
| if candidate_y_clean: |
| candidate_y_clean_np = np.concatenate(candidate_y_clean) |
| candidate_bins_clean_np = np.concatenate(candidate_bins_clean) if candidate_bins_clean else None |
| add_candidate_metrics( |
| metrics, |
| f"{prefix}_clean", |
| clean_true, |
| candidate_y_clean_np, |
| candidate_bins_clean_np, |
| z_bins=model.z_bins, |
| y_min=model.y_min, |
| y_max=model.y_max, |
| ) |
| metrics[f"{prefix}_clean/z_count"] = float(len(clean_true)) |
| if first_batch is not None and first_pred is not None: |
| if first_rec is not None and "target_flux" in first_batch and "loss_mask" in first_batch: |
| plot_reconstruction_batch( |
| run_dir / "plots" / f"{prefix}_reconstruction_step_{step:06d}.png", |
| first_batch["loglam"].numpy(), |
| first_batch["target_flux"].numpy(), |
| first_rec, |
| first_batch["loss_mask"].numpy(), |
| first_batch["valid"].numpy(), |
| first_batch["z"].numpy(), |
| np.expm1(first_pred), |
| ) |
| plot_spectra_batch(run_dir / "plots" / f"{prefix}_spectra_step_{step:06d}.png", first_batch, first_pred) |
| model.train() |
| return metrics |
|
|
|
|
| def make_loader( |
| samples: list[dict[str, Any]], |
| indices: np.ndarray, |
| cfg: RawCollatorConfig, |
| args: argparse.Namespace, |
| train: bool, |
| sampler: WeightedRandomSampler | None = None, |
| ) -> DataLoader: |
| return DataLoader( |
| SpectraListDataset(samples, indices), |
| batch_size=args.batch_size, |
| shuffle=train and sampler is None, |
| sampler=sampler, |
| num_workers=args.num_workers, |
| pin_memory=True, |
| collate_fn=RawSpectraCollator(cfg, train=train, seed=args.seed + (0 if train else 1000)), |
| ) |
|
|
|
|
| def checkpoint_score( |
| mode: str, |
| val_metrics: dict[str, float], |
| ood_metrics: dict[str, float] | None, |
| z_alpha: float = 0.6, |
| desi_mae_ceiling: float = 0.0, |
| desi_mae_penalty: float = 0.0, |
| ) -> float: |
| def score_prefix(metrics: dict[str, float], prefix: str) -> float: |
| z_score = ( |
| metrics.get(f"{prefix}/nmad", math.inf) |
| + metrics.get(f"{prefix}/cat_0p01", 1.0) |
| + metrics.get(f"{prefix}/mae_log1p", 1.0) |
| ) |
| rec_score = metrics.get(f"{prefix}/rec") |
| if rec_score is None or not math.isfinite(float(rec_score)): |
| return z_score |
| alpha = min(max(float(z_alpha), 0.0), 1.0) |
| return alpha * z_score + (1.0 - alpha) * float(rec_score) |
|
|
| val_score = score_prefix(val_metrics, "val") |
| if mode == "rec": |
| return float(val_metrics.get("val/rec", math.inf)) |
| if mode == "val" or ood_metrics is None: |
| score = val_score |
| if desi_mae_ceiling > 0 and desi_mae_penalty > 0: |
| val_mae = float(val_metrics.get("val/mae_z", 0.0)) |
| score += float(desi_mae_penalty) * max(0.0, val_mae - float(desi_mae_ceiling)) |
| return score |
| ood_score = score_prefix(ood_metrics, "ood") |
| if mode == "ood": |
| score = ood_score |
| else: |
| score = 0.5 * val_score + 0.5 * ood_score |
| if desi_mae_ceiling > 0 and desi_mae_penalty > 0: |
| val_mae = float(val_metrics.get("val/mae_z", 0.0)) |
| score += float(desi_mae_penalty) * max(0.0, val_mae - float(desi_mae_ceiling)) |
| return score |
|
|
|
|
| def scheduled_lr(base_lr: float, min_lr: float, step: int, total_steps: int, warmup_steps: int) -> float: |
| if warmup_steps > 0 and step <= warmup_steps: |
| return base_lr * float(step) / float(max(1, warmup_steps)) |
| if min_lr < 0 or total_steps <= warmup_steps: |
| return base_lr |
| progress = (step - warmup_steps) / float(max(1, total_steps - warmup_steps)) |
| progress = min(max(progress, 0.0), 1.0) |
| return min_lr + 0.5 * (base_lr - min_lr) * (1.0 + math.cos(math.pi * progress)) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--dataset-name", default="MultimodalUniverse/desi") |
| parser.add_argument("--max-samples", type=int, default=4096) |
| parser.add_argument("--cache-dir", default="/workspace/native_specz_mae/cache") |
| parser.add_argument("--hf-cache-dir", default=os.environ.get("HF_DATASETS_CACHE", "/workspace/hf_cache/datasets")) |
| parser.add_argument("--run-dir", default="/workspace/runs/hybrid_specz") |
| parser.add_argument("--resume-checkpoint", default="") |
| parser.add_argument("--allow-mismatched-checkpoint", action="store_true") |
| parser.add_argument("--refresh-data", action="store_true") |
| parser.add_argument("--epochs", type=int, default=8) |
| parser.add_argument("--batch-size", type=int, default=64) |
| parser.add_argument("--num-workers", type=int, default=2) |
| parser.add_argument("--target-length", type=int, default=4096) |
| parser.add_argument("--architecture", choices=["conv", "wave_token"], default="conv") |
| parser.add_argument("--d-model", type=int, default=256) |
| parser.add_argument("--conv-width", type=int, default=128) |
| parser.add_argument("--layers", type=int, default=5) |
| parser.add_argument("--heads", type=int, default=8) |
| parser.add_argument("--dropout", type=float, default=0.1) |
| parser.add_argument("--z-bins", type=int, default=64) |
| parser.add_argument("--stem-stride", type=int, choices=[4, 8], default=8) |
| parser.add_argument("--token-stride", type=int, default=8) |
| parser.add_argument("--rec-hidden-mult", type=int, default=0) |
| parser.add_argument("--rec-refine-width", type=int, default=16) |
| parser.add_argument("--rec-refine-kernel", type=int, default=5) |
| parser.add_argument("--layerscale-init", type=float, default=0.0) |
| parser.add_argument( |
| "--prediction-mode", |
| choices=[ |
| "regression", |
| "softbin", |
| "hybrid", |
| "bin_residual", |
| "ranked_bin_residual", |
| "candidate_rerank", |
| "calibrated_bin_residual", |
| ], |
| default="regression", |
| ) |
| parser.add_argument("--bin-temperature", type=float, default=1.0) |
| parser.add_argument("--residual-scale", type=float, default=0.06) |
| parser.add_argument("--candidate-topk", type=int, default=5) |
| parser.add_argument("--lr", type=float, default=2e-4) |
| parser.add_argument("--min-lr", type=float, default=-1.0) |
| parser.add_argument("--warmup-steps", type=int, default=0) |
| parser.add_argument("--weight-decay", type=float, default=0.03) |
| parser.add_argument("--grad-clip", type=float, default=1.0) |
| parser.add_argument("--grad-accum-steps", type=int, default=1) |
| parser.add_argument("--eval-every", type=int, default=100) |
| parser.add_argument("--eval-max-val", type=int, default=800) |
| parser.add_argument("--eval-max-ood", type=int, default=480) |
| parser.add_argument("--max-steps", type=int, default=0) |
| parser.add_argument("--checkpoint-score", choices=["val", "ood", "combined", "rec"], default="combined") |
| parser.add_argument("--score-z-alpha", type=float, default=0.6) |
| parser.add_argument("--desi-mae-ceiling", type=float, default=0.0) |
| parser.add_argument("--desi-mae-penalty", type=float, default=0.0) |
| parser.add_argument("--objective", choices=["joint", "rec_only", "z_only"], default="joint") |
| parser.add_argument("--freeze-mode", choices=["none", "adapter", "rerank", "calib"], default="none") |
| parser.add_argument("--train-top-layers", type=int, default=0) |
| parser.add_argument("--train-layernorms", action="store_true") |
| parser.add_argument("--replay-checkpoint", default="") |
| parser.add_argument("--replay-y-weight", type=float, default=0.0) |
| parser.add_argument("--replay-bin-weight", type=float, default=0.0) |
| parser.add_argument("--replay-clean-only", action="store_true") |
| parser.add_argument("--balance-redshift", action="store_true") |
| parser.add_argument("--train-clean-only", action="store_true") |
| parser.add_argument("--clean-sample-boost", type=float, default=1.0) |
| parser.add_argument("--augment-ood", action="store_true") |
| parser.add_argument("--eval-ood", action="store_true") |
| parser.add_argument("--random-mask-ratio", type=float, default=0.0) |
| parser.add_argument("--eval-mask-ratio", type=float, default=0.25) |
| parser.add_argument("--mask-mode", choices=["pixel", "span", "line_span", "mixed_span"], default="pixel") |
| parser.add_argument("--mask-span-min", type=int, default=16) |
| parser.add_argument("--mask-span-max", type=int, default=64) |
| parser.add_argument("--line-region-percentile", type=float, default=90.0) |
| parser.add_argument("--crop-prob", type=float, default=0.0) |
| parser.add_argument("--bad-window-prob", type=float, default=0.0) |
| parser.add_argument("--throughput-prob", type=float, default=0.0) |
| parser.add_argument("--noise-prob", type=float, default=0.0) |
| parser.add_argument("--resolution-prob", type=float, default=0.0) |
| parser.add_argument("--downsample-prob", type=float, default=0.0) |
| parser.add_argument("--line-dropout-prob", type=float, default=0.0) |
| parser.add_argument("--span-dropout-prob", type=float, default=0.0) |
| parser.add_argument("--grid-jitter-prob", type=float, default=0.0) |
| parser.add_argument("--grid-shift-frac", type=float, default=0.0) |
| parser.add_argument("--grid-scale-frac", type=float, default=0.0) |
| parser.add_argument("--grid-jitter-warmup-steps", type=int, default=0) |
| parser.add_argument("--redshift-shift", type=float, default=0.0) |
| parser.add_argument("--rec-weight", type=float, default=0.0) |
| parser.add_argument("--z-weight", type=float, default=1.0) |
| parser.add_argument("--z-bin-weight", type=float, default=0.25) |
| parser.add_argument("--z-candidate-weight", type=float, default=0.0) |
| parser.add_argument("--z-rerank-weight", type=float, default=0.0) |
| parser.add_argument("--z-nll-weight", type=float, default=0.05) |
| parser.add_argument("--zwarn-weight", type=float, default=0.3) |
| parser.add_argument("--high-z-boost", type=float, default=1.0) |
| parser.add_argument("--high-z-threshold", type=float, default=1.0) |
| parser.add_argument("--clean-z-only", action="store_true") |
| parser.add_argument("--seed", type=int, default=17) |
| args = parser.parse_args() |
|
|
| torch.manual_seed(args.seed) |
| np.random.seed(args.seed) |
| run_dir = Path(args.run_dir) / time.strftime("%Y%m%d_%H%M%S") |
| run_dir.mkdir(parents=True, exist_ok=True) |
| (run_dir / "args.json").write_text(json.dumps(vars(args), indent=2), encoding="utf-8") |
|
|
| samples = collect_mmu_desi( |
| cache_file=Path(args.cache_dir) / f"desi_{args.max_samples}.pt", |
| max_samples=args.max_samples, |
| dataset_name=args.dataset_name, |
| hf_cache_dir=args.hf_cache_dir, |
| refresh=args.refresh_data, |
| ) |
| stats = compute_sample_stats(samples) |
| (run_dir / "data_stats.json").write_text(json.dumps(stats.__dict__, indent=2), encoding="utf-8") |
| print("DATA_STATS", json.dumps(stats.__dict__, sort_keys=True)) |
|
|
| train_idx, val_idx = split_indices(len(samples), val_fraction=0.15, seed=args.seed) |
| if args.train_clean_only: |
| clean_train = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_) |
| train_idx = train_idx[clean_train] |
| if len(train_idx) == 0: |
| raise RuntimeError("No clean ZWARN==0 samples are available for --train-clean-only.") |
| print(f"TRAIN_CLEAN_ONLY n_train={len(train_idx)}") |
|
|
| sampler = None |
| if args.balance_redshift or args.clean_sample_boost != 1.0: |
| weights = np.ones(len(train_idx), dtype=np.float32) |
| y_train = np.asarray([np.log1p(float(samples[int(i)]["z"])) for i in train_idx], dtype=np.float32) |
| if args.balance_redshift: |
| bins = np.linspace(float(y_train.min()), float(y_train.max()) + 1e-6, 28) |
| bin_id = np.clip(np.digitize(y_train, bins) - 1, 0, len(bins) - 2) |
| counts = np.bincount(bin_id, minlength=len(bins) - 1).astype(np.float32) |
| weights *= 1.0 / np.maximum(counts[bin_id], 1.0) |
| if args.clean_sample_boost != 1.0: |
| clean = np.asarray([not bool(samples[int(i)]["zwarn"]) for i in train_idx], dtype=np.bool_) |
| weights *= np.where(clean, float(args.clean_sample_boost), 1.0).astype(np.float32) |
| weights = weights / weights.mean() |
| sampler = WeightedRandomSampler(torch.as_tensor(weights, dtype=torch.double), num_samples=len(weights), replacement=True) |
|
|
| train_cfg = RawCollatorConfig( |
| target_length=args.target_length, |
| random_mask_ratio=args.random_mask_ratio, |
| eval_mask_ratio=args.eval_mask_ratio, |
| mask_mode=args.mask_mode, |
| mask_span_min=args.mask_span_min, |
| mask_span_max=args.mask_span_max, |
| line_region_percentile=args.line_region_percentile, |
| augment_ood=args.augment_ood, |
| crop_prob=args.crop_prob, |
| bad_window_prob=args.bad_window_prob, |
| throughput_prob=args.throughput_prob, |
| noise_prob=args.noise_prob, |
| resolution_prob=args.resolution_prob, |
| downsample_prob=args.downsample_prob, |
| line_dropout_prob=args.line_dropout_prob, |
| span_dropout_prob=args.span_dropout_prob, |
| grid_jitter_prob=args.grid_jitter_prob, |
| grid_shift_frac=args.grid_shift_frac, |
| grid_scale_frac=args.grid_scale_frac, |
| grid_jitter_warmup_steps=args.grid_jitter_warmup_steps, |
| redshift_shift=args.redshift_shift, |
| ) |
| val_cfg = RawCollatorConfig( |
| target_length=args.target_length, |
| eval_mask_ratio=args.eval_mask_ratio, |
| mask_mode=args.mask_mode, |
| mask_span_min=args.mask_span_min, |
| mask_span_max=args.mask_span_max, |
| line_region_percentile=args.line_region_percentile, |
| ) |
| ood_cfg = RawCollatorConfig( |
| target_length=args.target_length, |
| eval_mask_ratio=args.eval_mask_ratio, |
| mask_mode=args.mask_mode, |
| mask_span_min=args.mask_span_min, |
| mask_span_max=args.mask_span_max, |
| line_region_percentile=args.line_region_percentile, |
| augment_ood=True, |
| crop_prob=0.65, |
| bad_window_prob=0.45, |
| throughput_prob=0.65, |
| noise_prob=0.35, |
| resolution_prob=0.45, |
| downsample_prob=0.35, |
| grid_jitter_prob=0.65, |
| grid_shift_frac=0.04, |
| grid_scale_frac=0.08, |
| ) |
| train_loader = make_loader(samples, train_idx, train_cfg, args, train=True, sampler=sampler) |
| val_loader = make_loader(samples, val_idx, val_cfg, args, train=False) |
| ood_loader = make_loader(samples, val_idx, ood_cfg, args, train=False) if args.eval_ood else None |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model_kwargs = dict( |
| d_model=args.d_model, |
| conv_width=args.conv_width, |
| layers=args.layers, |
| heads=args.heads, |
| dropout=args.dropout, |
| z_bins=args.z_bins, |
| rec_hidden_mult=args.rec_hidden_mult, |
| rec_refine_width=args.rec_refine_width, |
| rec_refine_kernel=args.rec_refine_kernel, |
| layerscale_init=args.layerscale_init, |
| prediction_mode=args.prediction_mode, |
| bin_temperature=args.bin_temperature, |
| residual_scale=args.residual_scale, |
| candidate_topk=args.candidate_topk, |
| ) |
| if args.architecture == "wave_token": |
| model = WavelengthTokenSpecZ(token_stride=args.token_stride, **model_kwargs).to(device) |
| else: |
| model = HybridSpecZ(stem_stride=args.stem_stride, **model_kwargs).to(device) |
| if args.resume_checkpoint: |
| ckpt = torch.load(args.resume_checkpoint, map_location=device, weights_only=False) |
| state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt |
| load_checkpoint_into_model(model, state, allow_mismatched=args.allow_mismatched_checkpoint) |
| print(f"RESUME_CHECKPOINT {args.resume_checkpoint}") |
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"MODEL_PARAMS {n_params}") |
|
|
| teacher_model = None |
| if args.replay_y_weight > 0 or args.replay_bin_weight > 0: |
| replay_path = args.replay_checkpoint or args.resume_checkpoint |
| if not replay_path: |
| raise RuntimeError("--replay-checkpoint or --resume-checkpoint is required when replay weights are nonzero.") |
| if args.architecture == "wave_token": |
| teacher_model = WavelengthTokenSpecZ(token_stride=args.token_stride, **model_kwargs).to(device) |
| else: |
| teacher_model = HybridSpecZ(stem_stride=args.stem_stride, **model_kwargs).to(device) |
| ckpt = torch.load(replay_path, map_location=device, weights_only=False) |
| state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt |
| load_checkpoint_into_model(teacher_model, state, allow_mismatched=args.allow_mismatched_checkpoint) |
| teacher_model.eval() |
| for param in teacher_model.parameters(): |
| param.requires_grad = False |
| print(f"REPLAY_TEACHER {replay_path}") |
|
|
| trainable_params = configure_trainable_parameters(model, args.freeze_mode, args.train_top_layers, args.train_layernorms) |
| print(f"TRAINABLE_PARAMS {trainable_params} freeze_mode={args.freeze_mode} train_top_layers={args.train_top_layers}") |
| opt_params = [p for p in model.parameters() if p.requires_grad] |
| if not opt_params: |
| raise RuntimeError("No trainable parameters remain after freeze configuration.") |
| optimizer = AdamW(opt_params, lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95)) |
| rec_weight = args.rec_weight |
| z_weight = args.z_weight |
| z_bin_weight = args.z_bin_weight |
| z_candidate_weight = args.z_candidate_weight |
| z_rerank_weight = args.z_rerank_weight |
| z_nll_weight = args.z_nll_weight |
| if args.objective == "rec_only": |
| rec_weight = rec_weight if rec_weight > 0 else 1.0 |
| z_weight = 0.0 |
| z_bin_weight = 0.0 |
| z_candidate_weight = 0.0 |
| z_rerank_weight = 0.0 |
| z_nll_weight = 0.0 |
| elif args.objective == "z_only": |
| rec_weight = 0.0 |
|
|
| loss_cfg = LossConfig( |
| rec_weight=rec_weight, |
| z_weight=z_weight, |
| z_bin_weight=z_bin_weight, |
| z_candidate_weight=z_candidate_weight, |
| z_rerank_weight=z_rerank_weight, |
| z_nll_weight=z_nll_weight, |
| zwarn_weight=args.zwarn_weight, |
| clean_z_only=args.clean_z_only, |
| high_z_boost=args.high_z_boost, |
| high_z_threshold=math.log1p(args.high_z_threshold), |
| ) |
| best_score = math.inf |
| global_step = 0 |
| micro_step = 0 |
| grad_accum_steps = max(1, int(args.grad_accum_steps)) |
| total_train_steps = args.max_steps if args.max_steps else int(math.ceil(len(train_loader) / grad_accum_steps) * args.epochs) |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| for epoch in range(args.epochs): |
| pbar = tqdm(train_loader, desc=f"hybrid epoch {epoch}") |
| for batch in pbar: |
| micro_step += 1 |
| batch = move_to_device(batch, device) |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): |
| out = model(batch["x"], batch["valid"], batch["loglam"]) |
| loss, parts = redshift_total_loss(model, out, batch, loss_cfg) |
| if teacher_model is not None: |
| with torch.no_grad(): |
| teacher_out = teacher_model(batch["x"], batch["valid"], batch["loglam"]) |
| replay, replay_parts = replay_loss( |
| out, |
| teacher_out, |
| batch, |
| y_weight=args.replay_y_weight, |
| bin_weight=args.replay_bin_weight, |
| clean_only=args.replay_clean_only, |
| ) |
| loss = loss + replay |
| parts = {**parts, "loss": parts["loss"] + replay.detach(), "replay": replay.detach(), **replay_parts} |
| (loss / grad_accum_steps).backward() |
| if micro_step % grad_accum_steps != 0: |
| pbar.set_postfix( |
| loss=float(parts["loss"].detach().cpu()), |
| rec=float(parts["rec"].detach().cpu()), |
| huber=float(parts["z_huber"].detach().cpu()), |
| replay=float(parts.get("replay", parts["loss"].sum() * 0.0).detach().cpu()), |
| accum=f"{micro_step % grad_accum_steps}/{grad_accum_steps}", |
| ) |
| continue |
| next_step = global_step + 1 |
| lr_now = scheduled_lr(args.lr, args.min_lr, next_step, total_train_steps, int(args.warmup_steps)) |
| for group in optimizer.param_groups: |
| group["lr"] = lr_now |
| grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| optimizer.step() |
| optimizer.zero_grad(set_to_none=True) |
| global_step = next_step |
| pbar.set_postfix( |
| loss=float(parts["loss"].detach().cpu()), |
| rec=float(parts["rec"].detach().cpu()), |
| huber=float(parts["z_huber"].detach().cpu()), |
| replay=float(parts.get("replay", parts["loss"].sum() * 0.0).detach().cpu()), |
| lr=lr_now, |
| grad=float(grad_norm.detach().cpu()) if torch.is_tensor(grad_norm) else float(grad_norm), |
| ) |
| if global_step == 1 or global_step % args.eval_every == 0: |
| val_metrics = evaluate( |
| model, |
| val_loader, |
| loss_cfg, |
| device, |
| run_dir, |
| global_step, |
| prefix="val", |
| max_batches=None, |
| max_examples=args.eval_max_val, |
| ) |
| print("VAL", global_step, json.dumps(val_metrics, sort_keys=True)) |
| ood_metrics = None |
| if ood_loader is not None: |
| ood_metrics = evaluate( |
| model, |
| ood_loader, |
| loss_cfg, |
| device, |
| run_dir, |
| global_step, |
| prefix="ood", |
| max_batches=None, |
| max_examples=args.eval_max_ood, |
| ) |
| print("OOD", global_step, json.dumps(ood_metrics, sort_keys=True)) |
| score = checkpoint_score( |
| args.checkpoint_score, |
| val_metrics, |
| ood_metrics, |
| z_alpha=args.score_z_alpha, |
| desi_mae_ceiling=args.desi_mae_ceiling, |
| desi_mae_penalty=args.desi_mae_penalty, |
| ) |
| if score < best_score: |
| best_score = score |
| best_metrics = {"step": global_step, "score": best_score, **val_metrics} |
| if ood_metrics is not None: |
| best_metrics.update(ood_metrics) |
| torch.save( |
| {"model": model.state_dict(), "args": vars(args), "step": global_step, "score": best_score, "metrics": best_metrics}, |
| run_dir / "best.pt", |
| ) |
| (run_dir / "best_metrics.json").write_text(json.dumps(best_metrics, indent=2), encoding="utf-8") |
| if args.max_steps and global_step >= args.max_steps: |
| break |
| if args.max_steps and global_step >= args.max_steps: |
| break |
|
|
| final_metrics = evaluate( |
| model, |
| val_loader, |
| loss_cfg, |
| device, |
| run_dir, |
| global_step, |
| prefix="val", |
| max_batches=None, |
| max_examples=args.eval_max_val, |
| ) |
| if ood_loader is not None: |
| final_metrics.update( |
| evaluate( |
| model, |
| ood_loader, |
| loss_cfg, |
| device, |
| run_dir, |
| global_step, |
| prefix="ood", |
| max_batches=None, |
| max_examples=args.eval_max_ood, |
| ) |
| ) |
| torch.save({"model": model.state_dict(), "args": vars(args), "step": global_step, "metrics": final_metrics}, run_dir / "last.pt") |
| (run_dir / "final_metrics.json").write_text(json.dumps(final_metrics, indent=2), encoding="utf-8") |
| print("FINAL", json.dumps(final_metrics, sort_keys=True)) |
| print("RUN_DIR", run_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
