| """Train CorrDiff's conditional mean, then its frozen-mean residual EDM.""" |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| from contextlib import nullcontext |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from torch import distributed as dist |
| from torch.nn import functional as F |
| from torch.nn.parallel import DistributedDataParallel as DDP |
| from torch.utils.data import DataLoader, DistributedSampler, TensorDataset |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.corrdiff import CorrDiff |
|
|
|
|
| def scalar(archive, key, default=None): |
| if key not in archive: |
| if default is not None: |
| return default |
| raise ValueError(f"NPZ is missing required metadata: {key}") |
| value = archive[key] |
| if value.ndim != 0: |
| raise ValueError(f"NPZ metadata {key} must be a scalar") |
| return str(value.item()) |
|
|
|
|
| def reduced_average(total, count, device, distributed): |
| values = torch.tensor([total, count], dtype=torch.float64, device=device) |
| if distributed: |
| dist.all_reduce(values, op=dist.ReduceOp.SUM) |
| if values[1].item() == 0: |
| raise RuntimeError("Training stage processed no batches") |
| return (values[0] / values[1]).item() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(ROOT / "conf/config.yaml")) |
| args = parser.parse_args() |
| config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) |
| global_rank = int(os.getenv("RANK", 0)) |
| local_rank = int(os.getenv("LOCAL_RANK", 0)) |
| world = int(os.getenv("WORLD_SIZE", 1)) |
| distributed = world > 1 |
| if distributed: |
| dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") |
| use_cuda = torch.cuda.is_available() and config["runtime"]["device"] != "cpu" |
| device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") |
| if use_cuda: |
| torch.cuda.set_device(local_rank) |
| seed = config["seed"] + global_rank |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
| archive = np.load(ROOT / config["data"]["path"]) |
| protocol = scalar(archive, "protocol") |
| data_source = scalar(archive, "data_source") |
| if protocol != config["data"]["protocol"]: |
| raise ValueError(f"Expected protocol {config['data']['protocol']}, got {protocol}") |
| if not data_source: |
| raise ValueError("data_source must be a non-empty scalar") |
| coarse = torch.from_numpy(archive[config["data"]["input_key"]]) |
| target = torch.from_numpy(archive[config["data"]["target_key"]]) |
| if coarse.ndim != 4 or target.ndim != 4: |
| raise ValueError("CorrDiff input and target must be NCHW tensors") |
| if len(coarse) != len(target) or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]): |
| raise ValueError("NPZ tensor shapes do not match config") |
| dataset = TensorDataset(coarse, target) |
| sampler = DistributedSampler(dataset, shuffle=True) if distributed else None |
| loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, |
| shuffle=sampler is None, num_workers=config["training"]["num_workers"]) |
| model = CorrDiff(**config["model"]).to(device) |
| if distributed: |
| model = DDP(model, device_ids=[local_rank] if use_cuda else None, |
| find_unused_parameters=True) |
| base = model.module if distributed else model |
| reg_opt = torch.optim.AdamW(base.regression.parameters(), lr=config["training"]["learning_rate"]) |
| diff_opt = torch.optim.AdamW(base.diffusion.parameters(), lr=config["training"]["learning_rate"]) |
| amp = bool(config["training"]["amp"] and use_cuda) |
| scaler = torch.amp.GradScaler("cuda", enabled=amp) |
| autocast = (lambda: torch.amp.autocast("cuda", enabled=True)) if amp else nullcontext |
| history = [] |
|
|
| |
| for epoch in range(config["training"]["regression_epochs"]): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| model.train() |
| total = count = 0 |
| for batch_index, (coarse_batch, target_batch) in enumerate(loader): |
| coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device) |
| reg_opt.zero_grad(set_to_none=True) |
| with autocast(): |
| loss = F.mse_loss(model(coarse_batch, mode="mean"), target_batch) |
| scaler.scale(loss).backward() |
| scaler.step(reg_opt) |
| scaler.update() |
| total += loss.item() |
| count += 1 |
| if batch_index + 1 >= config["training"]["max_batches_per_epoch"]: |
| break |
| value = reduced_average(total, count, device, distributed) |
| record = {"stage": "regression", "epoch": epoch + 1, "regression_mse": value} |
| history.append(record) |
| if global_rank == 0: |
| print(json.dumps(record)) |
|
|
| base.regression.eval() |
| for parameter in base.regression.parameters(): |
| parameter.requires_grad_(False) |
| for epoch in range(config["training"]["diffusion_epochs"]): |
| if sampler is not None: |
| sampler.set_epoch(config["training"]["regression_epochs"] + epoch) |
| base.diffusion.train() |
| total = count = 0 |
| for batch_index, (coarse_batch, target_batch) in enumerate(loader): |
| coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device) |
| with torch.no_grad(): |
| mean = base.mean(coarse_batch) |
| residual = target_batch - mean |
| sigma = (torch.randn(len(coarse_batch), device=device) * config["training"]["p_std"] + config["training"]["p_mean"]).exp() |
| noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual) |
| diff_opt.zero_grad(set_to_none=True) |
| with autocast(): |
| denoised = model(coarse_batch, mode="denoise", mean=mean, noisy=noisy, sigma=sigma) |
| weight = (sigma.square() + base.sigma_data**2) / (sigma * base.sigma_data).square() |
| loss = (weight[:, None, None, None] * (denoised - residual).square()).mean() |
| scaler.scale(loss).backward() |
| scaler.step(diff_opt) |
| scaler.update() |
| total += loss.item() |
| count += 1 |
| if batch_index + 1 >= config["training"]["max_batches_per_epoch"]: |
| break |
| value = reduced_average(total, count, device, distributed) |
| record = {"stage": "diffusion", "epoch": epoch + 1, "edm_loss": value} |
| history.append(record) |
| if global_rank == 0: |
| print(json.dumps(record)) |
|
|
| if global_rank == 0: |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"model": base.state_dict(), "config": config, "format": "corrdiff-edm-v3", |
| "protocol": protocol, "data_source": data_source}, checkpoint) |
| metrics = ROOT / config["paths"]["training_metrics"] |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"history": history, "protocol": protocol, |
| "data_source": data_source}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint}") |
| if distributed: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
