File size: 7,677 Bytes
1d4cac8 | 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 167 168 | """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()
|