shaoxiongduan
92cc893cfc
Training:
- finetune.py: optional track/base param groups (WANTRACK_TRACK_GROUP) so the
track pathway can take a different LR (WANTRACK_TRACK_LR_MULT); patch_embedding
gradient masked to the track slot [:, 36:]. Raises if nothing matches.
- track_warmup.py callback: hold the pretrained DiT at lr 0 while the track
pathway warms up, then ramp it back in.
14B configs + held-alloc launchers:
- finetune_wantrack_{synth_sparse_fixed,synth_sparse_random,openvid_sparse_merged_bias,
synth_stage2}_14b + openvid_14b_{warmup,lwlr} + 1.3B stage2 configs.
- run_wan14b_held.sh (held alloc + auto-restart + rank cleanup, per-stage WANTRACK
env), acquire_nodes.sh, chain_stepAB.sh, run_stepB_seed.sh, run_stepC_merge.sh,
watch_gate.sh, run_synth_stage2_slurm.sh.
Init / diagnostics:
- convert_trackwan_init_v2.py: --pe-src to lift patch_embedding[:, 36:] and
--track-src to lift track_encoder together (co-adapted merge); sharded-source support.
- check_gate_growth.py, gradient_flow_analysis.py, gradient_flow_summarize.py.
Data pipeline: segment_tracks.py sharding; decode/segment/tracks slurm scripts;
precompute_null_text, two_track_val, cfg_ablation, lwlr_experiment.
Inference: app_multi_trace.py updates.
471 lines
20 KiB
Python
471 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Per-sample gradient / SNR analysis of the WanTrack DiT, comparing 4 init/checkpoint variants
|
|
(zero_init, random_init, merged, overfit_full) on real OpenVid parquet samples.
|
|
|
|
For each (model, layer) we record:
|
|
avg_per_sample_norm = mean_i ||g_i|| (typical bs=1 grad magnitude)
|
|
mean_grad_norm = ||mean_i g_i|| (grad magnitude after bs=N averaging)
|
|
snr = mean_grad_norm / avg_per_sample_norm
|
|
|
|
High SNR => signal survives averaging (universal component).
|
|
Low SNR => signal cancels (sample-specific / noise).
|
|
|
|
Layers tracked (custom-model names, not safetensors names):
|
|
patch_embedding.proj.weight [:, :36] and [:, 36:]
|
|
track_encoder.proj.weight
|
|
track_encoder.temporal_conv.weight
|
|
blocks.0.to_q.weight (early self-attn Q)
|
|
blocks.14.ffn.fc_in.weight (mid FFN in-proj)
|
|
blocks.29.to_out.weight (late self-attn out)
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import random
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow.parquet as pq
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors.torch import load_file as st_load
|
|
|
|
# Init distributed BEFORE any fastvideo import that touches SP world.
|
|
def _init_distributed_1gpu() -> None:
|
|
os.environ.setdefault("RANK", "0")
|
|
os.environ.setdefault("LOCAL_RANK", "0")
|
|
os.environ.setdefault("WORLD_SIZE", "1")
|
|
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
|
os.environ.setdefault("MASTER_PORT", "29511")
|
|
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
from fastvideo.distributed.parallel_state import (
|
|
maybe_init_distributed_environment_and_model_parallel,
|
|
)
|
|
maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
|
|
|
|
|
|
_init_distributed_1gpu()
|
|
|
|
# Now safe to import model classes
|
|
from fastvideo.configs.models.dits.trackwan import ( # noqa: E402
|
|
TrackWanVideoArchConfig,
|
|
TrackWanVideoConfig,
|
|
)
|
|
from fastvideo.forward_context import set_forward_context # noqa: E402
|
|
from fastvideo.models.dits.trackwan.model import TrackWanTransformer3DModel # noqa: E402
|
|
from fastvideo.models.loader.utils import ( # noqa: E402
|
|
get_param_names_mapping,
|
|
hf_to_custom_state_dict,
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Model set-up
|
|
# ------------------------------------------------------------------ #
|
|
MODELS = {
|
|
"zero_init": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_nobias_init",
|
|
"random_init": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_nobias_random_init",
|
|
"merged": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_merged_from_overfit",
|
|
"overfit_full": "/mnt/lustre/vlm-s4duan/models/overfit_2000",
|
|
"fun_control": "/mnt/lustre/vlm-s4duan/models/trackwan_1.3b_i2v_d64_from_fun_control_init",
|
|
# 14B variants
|
|
"14b_zero_init": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_nobias_init",
|
|
"14b_random_init": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_random_init",
|
|
"14b_partial_merged": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_partial_merged",
|
|
# zero-gate init WITH bias (the base the step-A overfit started from — gate exactly 0).
|
|
"14b_zero_init_bias": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_zero_init_bias",
|
|
# fresh 14B base with the FIXED-overfit (step A, ckpt-2000) encoder + gate merged in — this
|
|
# is the candidate stage-1 init; its per-sample track grad decides whether it will train.
|
|
"14b_merged_fixed2000": "/mnt/lustre/vlm-s4duan/models/trackwan_14b_i2v_d64_merged_from_fixed2000_bias",
|
|
}
|
|
|
|
# Custom-model parameter names we snapshot per sample.
|
|
# Special slicing for patch_embedding is handled below.
|
|
TRACK_LAYERS = [
|
|
"patch_embedding.proj.weight", # will split into base=[:, :36] / track=[:, 36:]
|
|
"track_encoder.proj.weight",
|
|
"track_encoder.temporal_conv.weight",
|
|
"blocks.0.to_q.weight",
|
|
"blocks.14.ffn.fc_in.weight",
|
|
"blocks.29.to_out.weight",
|
|
]
|
|
|
|
|
|
def build_model_from_diffusers(model_dir: str, device: torch.device) -> TrackWanTransformer3DModel:
|
|
cfg_path = Path(model_dir) / "transformer" / "config.json"
|
|
with open(cfg_path, "r") as f:
|
|
hf_config = json.load(f)
|
|
|
|
arch = TrackWanVideoArchConfig(
|
|
num_attention_heads=int(hf_config["num_attention_heads"]),
|
|
attention_head_dim=int(hf_config["attention_head_dim"]),
|
|
in_channels=int(hf_config["in_channels"]),
|
|
out_channels=int(hf_config["out_channels"]),
|
|
text_dim=int(hf_config.get("text_dim", 4096)),
|
|
freq_dim=int(hf_config.get("freq_dim", 256)),
|
|
ffn_dim=int(hf_config["ffn_dim"]),
|
|
num_layers=int(hf_config["num_layers"]),
|
|
cross_attn_norm=bool(hf_config.get("cross_attn_norm", True)),
|
|
qk_norm=str(hf_config.get("qk_norm", "rms_norm_across_heads")),
|
|
eps=float(hf_config.get("eps", 1e-6)),
|
|
image_dim=int(hf_config.get("image_dim") or 0) or None,
|
|
added_kv_proj_dim=int(hf_config.get("added_kv_proj_dim") or 0) or None,
|
|
rope_max_seq_len=int(hf_config.get("rope_max_seq_len", 1024)),
|
|
track_config=dict(hf_config.get("track_config", {})),
|
|
)
|
|
# The trackwan config has an "in_channels: 52" default, but we still pass hf_config so
|
|
# TrackWanTransformer3DModel picks up the correct track_channels from track_config.
|
|
cfg = TrackWanVideoConfig(arch_config=arch)
|
|
|
|
# Ensure downstream consumers see the right in_channels regardless of default
|
|
cfg.arch_config.in_channels = int(hf_config["in_channels"])
|
|
cfg.arch_config.out_channels = int(hf_config["out_channels"])
|
|
|
|
model = TrackWanTransformer3DModel(config=cfg, hf_config=hf_config)
|
|
model = model.to(device=device, dtype=torch.bfloat16)
|
|
|
|
# --- load weights ---
|
|
st_files = sorted(glob.glob(str(Path(model_dir) / "transformer" / "*.safetensors")))
|
|
assert len(st_files) >= 1, f"no safetensors under {model_dir}/transformer"
|
|
hf_sd: dict[str, torch.Tensor] = {}
|
|
for p in st_files:
|
|
hf_sd.update(st_load(p))
|
|
|
|
mapping_fn = get_param_names_mapping(model.param_names_mapping)
|
|
custom_sd, _ = hf_to_custom_state_dict(hf_sd, mapping_fn)
|
|
|
|
# Load into model
|
|
model_sd = model.state_dict()
|
|
missing, unexpected = [], []
|
|
for k, v in custom_sd.items():
|
|
if k in model_sd:
|
|
model_sd[k].data.copy_(v.to(model_sd[k].dtype))
|
|
else:
|
|
unexpected.append(k)
|
|
for k in model_sd:
|
|
if k not in custom_sd:
|
|
missing.append(k)
|
|
print(f"[{model_dir}] missing={len(missing)} unexpected={len(unexpected)}")
|
|
if unexpected:
|
|
print(" unexpected head:", unexpected[:5])
|
|
if missing:
|
|
print(" missing head:", missing[:5])
|
|
|
|
# Gradient checkpointing needs the forward context to be re-established during
|
|
# backward recompute; we keep both fwd+bwd wrapped in set_forward_context in
|
|
# run_one_model, so grad_checkpointing is safe.
|
|
model.gradient_checkpointing = True
|
|
model._gradient_checkpointing_func = lambda fn, *args, **kwargs: torch.utils.checkpoint.checkpoint(
|
|
fn, *args, use_reentrant=False, **kwargs
|
|
)
|
|
|
|
# FastVideo's linear layers create Parameters with requires_grad=False (FSDP path).
|
|
# For gradient analysis we force all params trainable.
|
|
for p in model.parameters():
|
|
p.requires_grad_(True)
|
|
|
|
model.train() # ensure gradients flow
|
|
return model
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Parquet loader
|
|
# ------------------------------------------------------------------ #
|
|
@dataclass
|
|
class Sample:
|
|
vae_latent: torch.Tensor # [16, T, H, W]
|
|
first_frame_latent: torch.Tensor # [16, T, H, W]
|
|
clip_feature: torch.Tensor # [257, 1280]
|
|
text_embedding: torch.Tensor # [L, 4096]
|
|
track_points: torch.Tensor # [T_full, N, 2]
|
|
track_visibility: torch.Tensor # [T_full, N]
|
|
object_ids: torch.Tensor # [N]
|
|
track_weights: torch.Tensor # [N]
|
|
|
|
|
|
def _decode(row, name) -> torch.Tensor:
|
|
shape = row[f"{name}_shape"]
|
|
data = row[f"{name}_bytes"]
|
|
arr = np.frombuffer(data, dtype=np.float32).reshape(shape).copy()
|
|
return torch.from_numpy(arr)
|
|
|
|
|
|
def load_samples(pattern: str, n: int, seed: int = 0) -> list[Sample]:
|
|
files = sorted(glob.glob(pattern, recursive=True))
|
|
random.Random(seed).shuffle(files)
|
|
samples: list[Sample] = []
|
|
for f in files:
|
|
if len(samples) >= n:
|
|
break
|
|
try:
|
|
tbl = pq.read_table(f)
|
|
except Exception as e:
|
|
print("skip", f, e)
|
|
continue
|
|
for row in tbl.to_pylist():
|
|
try:
|
|
s = Sample(
|
|
vae_latent=_decode(row, "vae_latent"),
|
|
first_frame_latent=_decode(row, "first_frame_latent"),
|
|
clip_feature=_decode(row, "clip_feature"),
|
|
text_embedding=_decode(row, "text_embedding"),
|
|
track_points=_decode(row, "track_points"),
|
|
track_visibility=_decode(row, "track_visibility"),
|
|
object_ids=_decode(row, "object_ids"),
|
|
track_weights=_decode(row, "track_weights"),
|
|
)
|
|
except Exception as e:
|
|
print("skip row", e)
|
|
continue
|
|
samples.append(s)
|
|
if len(samples) >= n:
|
|
break
|
|
print(f"loaded {len(samples)} samples from {pattern}")
|
|
return samples
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# I2V cond concat and flow-matching noise
|
|
# ------------------------------------------------------------------ #
|
|
def build_i2v_cond_concat(first_frame_latent: torch.Tensor,
|
|
vae_temporal_compression: int = 4) -> torch.Tensor:
|
|
"""[B,16,T,H,W] first-frame latent -> [B,20,T,H,W] (4 mask + 16 latent).
|
|
|
|
Mirrors WanTrackModel._build_i2v_cond_concat (mask=1 on first frame slot only).
|
|
"""
|
|
b, _, num_latent_t, h, w = first_frame_latent.shape
|
|
ratio = vae_temporal_compression
|
|
num_frames = (num_latent_t - 1) * ratio + 1
|
|
mask = torch.ones(b, 1, num_frames, h, w,
|
|
device=first_frame_latent.device, dtype=first_frame_latent.dtype)
|
|
mask[:, :, 1:] = 0
|
|
first = torch.repeat_interleave(mask[:, :, :1], dim=2, repeats=ratio)
|
|
mask = torch.cat([first, mask[:, :, 1:]], dim=2)
|
|
mask = mask.view(b, -1, ratio, h, w).transpose(1, 2) # [B, ratio, num_latent_t, H, W]
|
|
return torch.cat([mask, first_frame_latent], dim=1)
|
|
|
|
|
|
def flow_matching_noisy(clean: torch.Tensor, noise: torch.Tensor,
|
|
sigma: float, flow_shift: float = 6.0) -> tuple[torch.Tensor, float, float]:
|
|
"""(1-s)*x + s*n with shifted sigma. Returns noisy, shifted_sigma, timestep."""
|
|
shifted = flow_shift * sigma / (1 + (flow_shift - 1) * sigma)
|
|
noisy = (1.0 - shifted) * clean + shifted * noise
|
|
return noisy, shifted, shifted * 1000.0
|
|
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Per-sample gradient snapshot
|
|
# ------------------------------------------------------------------ #
|
|
def snapshot_layer_grads(model: torch.nn.Module) -> dict[str, torch.Tensor]:
|
|
out: dict[str, torch.Tensor] = {}
|
|
param_map = dict(model.named_parameters())
|
|
for name in TRACK_LAYERS:
|
|
if name not in param_map:
|
|
# Some models pack Q/K/V differently. Try nested attn1 name as fallback.
|
|
alt = name.replace("blocks.0.to_q", "blocks.0.attn1.to_q") \
|
|
.replace("blocks.29.to_out", "blocks.29.attn1.to_out") \
|
|
.replace(".ffn.fc_in.", ".ffn.net.0.proj.")
|
|
if alt in param_map:
|
|
p = param_map[alt]
|
|
else:
|
|
print(f"WARN param not found: {name}")
|
|
continue
|
|
else:
|
|
p = param_map[name]
|
|
if p.grad is None:
|
|
print(f"WARN no grad on {name}")
|
|
continue
|
|
out[name] = p.grad.detach().clone().float().cpu()
|
|
return out
|
|
|
|
|
|
def run_one_model(model_name: str, model_dir: str, samples: list[Sample],
|
|
out_dir: Path, device: torch.device,
|
|
flow_shift: float = 6.0) -> dict:
|
|
print(f"\n===== {model_name} =====")
|
|
torch.cuda.empty_cache()
|
|
model = build_model_from_diffusers(model_dir, device)
|
|
print("model params:",
|
|
sum(p.numel() for p in model.parameters()) / 1e9, "B")
|
|
|
|
# Store per-sample grads (may be large: patch_embedding.weight = 1536*52*4 = 320K params).
|
|
per_sample_grads: dict[str, list[torch.Tensor]] = {n: [] for n in TRACK_LAYERS}
|
|
|
|
for i, s in enumerate(samples):
|
|
model.zero_grad(set_to_none=True)
|
|
|
|
# Move tensors to device / bf16 dtype
|
|
vae_latent = s.vae_latent.unsqueeze(0).to(device, dtype=torch.bfloat16) # [1,16,T,H,W]
|
|
first_frame_latent = s.first_frame_latent.unsqueeze(0).to(device, dtype=torch.bfloat16)
|
|
clip_feature = s.clip_feature.unsqueeze(0).to(device, dtype=torch.bfloat16)
|
|
text_embedding = s.text_embedding.unsqueeze(0).to(device, dtype=torch.bfloat16)
|
|
# Note attention mask: use all ones over the raw text seq len
|
|
text_mask = torch.ones(1, s.text_embedding.shape[0], device=device, dtype=torch.bfloat16)
|
|
|
|
# Track tensors
|
|
track_points = s.track_points.unsqueeze(0).to(device, dtype=torch.bfloat16) # [1, T_full, N, 2]
|
|
track_visibility = s.track_visibility.unsqueeze(0).to(device, dtype=torch.bfloat16)
|
|
|
|
# Truncate track T to (num_latent_t - 1) * ratio + 1
|
|
num_latent_t = vae_latent.shape[2]
|
|
expected_T = (num_latent_t - 1) * 4 + 1
|
|
track_points = track_points[:, :expected_T]
|
|
track_visibility = track_visibility[:, :expected_T]
|
|
|
|
# Sample sigma ~ U(0,1)
|
|
sigma_u = float(torch.rand(1).item())
|
|
# Compute noise and noisy
|
|
noise = torch.randn_like(vae_latent, dtype=torch.float32)
|
|
clean_f = vae_latent.float()
|
|
noisy_f, shifted_sigma, timestep_val = flow_matching_noisy(clean_f, noise, sigma_u,
|
|
flow_shift=flow_shift)
|
|
noisy = noisy_f.to(torch.bfloat16) # [1, 16, T, H, W]
|
|
|
|
# Build I2V concat
|
|
cond20 = build_i2v_cond_concat(first_frame_latent, vae_temporal_compression=4) # [1,20,T,H,W]
|
|
hidden_states = torch.cat([noisy, cond20], dim=1) # [1, 36, T, H, W]
|
|
|
|
# Timestep tensor
|
|
timestep = torch.tensor([timestep_val], device=device, dtype=torch.bfloat16)
|
|
|
|
# Forward + backward BOTH inside set_forward_context so grad-checkpoint
|
|
# recompute finds the forward context.
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16), \
|
|
set_forward_context(current_timestep=timestep, attn_metadata=None):
|
|
pred = model(
|
|
hidden_states=hidden_states,
|
|
encoder_hidden_states=text_embedding,
|
|
encoder_attention_mask=text_mask,
|
|
timestep=timestep,
|
|
encoder_hidden_states_image=clip_feature,
|
|
track_points=track_points,
|
|
track_visibility=track_visibility,
|
|
return_dict=False,
|
|
)
|
|
# pred is [B, C, T, H, W]
|
|
# Loss: MSE(pred, noise - clean) computed in fp32
|
|
target = (noise - clean_f) # [B,16,T,H,W]
|
|
loss = F.mse_loss(pred.float(), target.float())
|
|
loss.backward()
|
|
|
|
# Snapshot
|
|
gs = snapshot_layer_grads(model)
|
|
for n, g in gs.items():
|
|
per_sample_grads[n].append(g)
|
|
|
|
if i < 3 or i % 8 == 0:
|
|
print(f" sample {i}: sigma_u={sigma_u:.3f} shifted={shifted_sigma:.3f} "
|
|
f"t={timestep_val:.1f} loss={loss.item():.4f} "
|
|
f"|g[patch_embed]|={gs.get('patch_embedding.proj.weight', torch.zeros(1)).norm().item():.4f}")
|
|
# free bookkeeping
|
|
del pred, hidden_states, noisy, noise, cond20, target, loss
|
|
torch.cuda.empty_cache()
|
|
|
|
# Save all grads
|
|
model_out = out_dir / model_name
|
|
model_out.mkdir(parents=True, exist_ok=True)
|
|
torch.save({n: torch.stack(gs, dim=0) for n, gs in per_sample_grads.items()
|
|
if len(gs) > 0},
|
|
model_out / "grads.pt")
|
|
|
|
# Compute stats
|
|
stats: dict = {}
|
|
for n in TRACK_LAYERS:
|
|
if n not in per_sample_grads or not per_sample_grads[n]:
|
|
continue
|
|
g_stack = torch.stack(per_sample_grads[n], dim=0) # [N, *param_shape]
|
|
# Special split for patch_embedding
|
|
if n == "patch_embedding.proj.weight":
|
|
# Shape [N, out, in, kD, kH, kW] with in=52
|
|
base_slice = g_stack[:, :, :36] # [N, out, 36, ...]
|
|
track_slice = g_stack[:, :, 36:] # [N, out, 16, ...]
|
|
for key, sub in [(f"{n}[:,:36]", base_slice), (f"{n}[:,36:]", track_slice)]:
|
|
per_sample_norms = sub.flatten(1).norm(dim=1) # [N]
|
|
mean_grad = sub.mean(0)
|
|
mean_grad_norm = mean_grad.flatten().norm().item()
|
|
avg_per_sample_norm = per_sample_norms.mean().item()
|
|
snr = mean_grad_norm / (avg_per_sample_norm + 1e-30)
|
|
stats[key] = {
|
|
"avg_per_sample_norm": avg_per_sample_norm,
|
|
"mean_grad_norm": mean_grad_norm,
|
|
"snr": snr,
|
|
"n_params": int(sub[0].numel()),
|
|
}
|
|
else:
|
|
per_sample_norms = g_stack.flatten(1).norm(dim=1)
|
|
mean_grad = g_stack.mean(0)
|
|
mean_grad_norm = mean_grad.flatten().norm().item()
|
|
avg_per_sample_norm = per_sample_norms.mean().item()
|
|
snr = mean_grad_norm / (avg_per_sample_norm + 1e-30)
|
|
stats[n] = {
|
|
"avg_per_sample_norm": avg_per_sample_norm,
|
|
"mean_grad_norm": mean_grad_norm,
|
|
"snr": snr,
|
|
"n_params": int(g_stack[0].numel()),
|
|
}
|
|
|
|
# Free model
|
|
del model
|
|
torch.cuda.empty_cache()
|
|
return stats
|
|
|
|
|
|
def print_table(all_stats: dict) -> None:
|
|
print("\n\n===== SNR TABLE =====")
|
|
print(f"{'model':<14} {'layer':<48} {'avg_per_sample':>16} {'mean_grad':>16} {'snr':>10}")
|
|
print("-" * 110)
|
|
for model, stats in all_stats.items():
|
|
for layer, s in stats.items():
|
|
print(f"{model:<14} {layer:<48} {s['avg_per_sample_norm']:>16.6f} "
|
|
f"{s['mean_grad_norm']:>16.6f} {s['snr']:>10.4f}")
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--n-samples", type=int, default=32)
|
|
ap.add_argument("--parquet-glob", type=str,
|
|
default="/mnt/lustre/vlm-s4duan/openvid_1m/combined_parquet_dataset/shard_*/**/*.parquet")
|
|
ap.add_argument("--out-dir", type=str, default="/mnt/lustre/vlm-s4duan/gradient_analysis")
|
|
ap.add_argument("--flow-shift", type=float, default=6.0)
|
|
ap.add_argument("--seed", type=int, default=42)
|
|
ap.add_argument("--only", type=str, default="", help="comma list of models to run")
|
|
args = ap.parse_args()
|
|
|
|
torch.manual_seed(args.seed)
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
|
|
out_dir = Path(args.out_dir)
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
device = torch.device("cuda")
|
|
|
|
samples = load_samples(args.parquet_glob, args.n_samples, seed=args.seed)
|
|
assert len(samples) == args.n_samples, "not enough samples loaded"
|
|
|
|
all_stats: dict = {}
|
|
models_to_run = args.only.split(",") if args.only else list(MODELS.keys())
|
|
for name in models_to_run:
|
|
name = name.strip()
|
|
if not name:
|
|
continue
|
|
stats = run_one_model(name, MODELS[name], samples, out_dir, device,
|
|
flow_shift=args.flow_shift)
|
|
all_stats[name] = stats
|
|
# Save partial after each model
|
|
with open(out_dir / "summary.json", "w") as f:
|
|
json.dump(all_stats, f, indent=2)
|
|
|
|
print_table(all_stats)
|
|
print(f"\nsummary -> {out_dir / 'summary.json'}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|