| """Train compact DOFA across strictly validated, variable-channel sensors.""" |
|
|
| import importlib.util |
| import json |
| import os |
| import random |
| 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.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_model_class(): |
| spec = importlib.util.spec_from_file_location("dofa_model", ROOT / "model" / "dofa.py") |
| module = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(module) |
| return module.DOFA |
|
|
|
|
| def expected_wavelengths(modality): |
| if modality.get("wavelength_mode") == "synthetic_uniform": |
| return modality["wavelength_start"] + np.arange(modality["channels"], dtype=np.float32) * modality["wavelength_step"] |
| return np.asarray(modality["wavelengths"], dtype=np.float32) |
|
|
|
|
| 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}") |
| if archive[key].ndim != 0: |
| raise ValueError(f"NPZ metadata {key} must be a scalar") |
| return archive[key].item() |
|
|
|
|
| class SensorDataset(Dataset): |
| def __init__(self, path, name, data_config): |
| archive = np.load(path) |
| if "images" not in archive or "wavelengths" not in archive: |
| raise ValueError(f"{path} must contain images and wavelengths") |
| self.images = archive["images"] |
| wavelengths = archive["wavelengths"] |
| modality = data_config["modalities"][name] |
| expected = expected_wavelengths(modality) |
| if self.images.ndim != 4: |
| raise ValueError(f"{path}: images must be NCHW") |
| if self.images.dtype != np.float32: |
| raise ValueError(f"{path}: images must use float32") |
| if self.images.shape[1:] != (modality["channels"], data_config["image_size"], data_config["image_size"]): |
| raise ValueError(f"{path}: image shape does not match configured channels/224x224") |
| if wavelengths.shape != (modality["channels"],) or not np.issubdtype(wavelengths.dtype, np.floating): |
| raise ValueError(f"{path}: wavelengths must be a floating [C] array") |
| if not np.isfinite(wavelengths).all() or not np.allclose(wavelengths, expected, rtol=1e-5, atol=1e-6): |
| raise ValueError(f"{path}: wavelengths do not match configured sensor wavelengths") |
| self.modality = str(scalar(archive, "modality")) |
| self.protocol = str(scalar(archive, "protocol")) |
| self.data_source = str(scalar(archive, "data_source", "unknown")) |
| if self.modality != name: |
| raise ValueError(f"{path}: modality {self.modality} does not match {name}") |
| if self.protocol != data_config["protocol"]: |
| raise ValueError(f"{path}: protocol {self.protocol} does not match config") |
| self.wavelengths = torch.from_numpy(wavelengths.astype("float32", copy=False)) |
|
|
| def __len__(self): |
| return len(self.images) |
|
|
| def __getitem__(self, index): |
| return torch.from_numpy(self.images[index]) |
|
|
|
|
| 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 processed no batches") |
| return (values[0] / values[1]).item() |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text(encoding="utf-8")) |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| global_rank = int(os.environ.get("RANK", "0")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| distributed = world_size > 1 |
| if distributed: |
| dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo") |
| requested = config["runtime"]["device"] |
| use_accelerator = torch.cuda.is_available() and requested != "cpu" |
| device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu") |
| if use_accelerator: |
| torch.cuda.set_device(local_rank) |
| seed = config["seed"] + global_rank |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
| data_root = ROOT / config["data"]["root"] |
| loaders = [] |
| for modality in config["data"]["modalities"]: |
| path = data_root / f"train_{modality}.npz" |
| if not path.exists(): |
| raise FileNotFoundError(f"Missing training data: {path.relative_to(ROOT)}. Run `python scripts/fake_data.py` first.") |
| dataset = SensorDataset(path, modality, config["data"]) |
| sampler = DistributedSampler(dataset, shuffle=True) if distributed else None |
| loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], |
| shuffle=sampler is None, sampler=sampler, |
| num_workers=config["training"]["num_workers"]) |
| loaders.append((dataset, loader, sampler)) |
|
|
| model = load_model_class()(**config["model"]).to(device) |
| if distributed: |
| model = DistributedDataParallel(model, device_ids=[local_rank] if use_accelerator else None) |
| optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"], |
| weight_decay=config["training"]["weight_decay"]) |
| amp = bool(config["training"]["amp"] and use_accelerator) |
| 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"]["epochs"]): |
| model.train() |
| records = {} |
| for dataset, loader, sampler in loaders: |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| total = count = 0 |
| for images in loader: |
| optimizer.zero_grad(set_to_none=True) |
| with autocast(): |
| output = model(images.to(device), dataset.wavelengths.to(device)) |
| scaler.scale(output["loss"]).backward() |
| scaler.step(optimizer) |
| scaler.update() |
| total += output["loss"].item() |
| count += 1 |
| records[dataset.modality] = reduced_average(total, count, device, distributed) |
| history.append({"epoch": epoch + 1, "reconstruction_loss": records}) |
| if global_rank == 0: |
| print(f"epoch={epoch + 1} " + " ".join(f"{name}={value:.6f}" for name, value in records.items())) |
|
|
| if global_rank == 0: |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| metrics = ROOT / config["paths"]["training_metrics"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| base_model = model.module if hasattr(model, "module") else model |
| sources = sorted({dataset.data_source for dataset, _, _ in loaders}) |
| torch.save({"model": base_model.state_dict(), "config": config, |
| "protocol": config["data"]["protocol"], "data_sources": sources}, checkpoint) |
| metrics.write_text(json.dumps({"history": history, |
| "modalities": list(config["data"]["modalities"]), |
| "protocol": config["data"]["protocol"], |
| "data_sources": sources}, indent=2) + "\n", encoding="utf-8") |
| print(f"checkpoint={checkpoint.relative_to(ROOT)}") |
| if distributed: |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
