File size: 7,512 Bytes
5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 0face05 5b3329b 0face05 5b3329b 0face05 5b3329b 0face05 5b3329b 986404c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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 = []
# Stage 1 is completed in full before any residual-EDM update occurs.
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()
|