| """Train DINCAE with random cloud masking and masked Gaussian NLL.""" |
|
|
| import json |
| import os |
| import sys |
| from datetime import date |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, Dataset, DistributedSampler |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| from model.dincae import DINCAE, build_input, masked_gaussian_nll |
|
|
|
|
| class SSTDataset(Dataset): |
| def __init__(self, arrays, indices, cloud_fraction, seed): |
| self.data, self.indices = arrays, list(indices) |
| self.cloud_fraction, self.seed = float(cloud_fraction), int(seed) |
|
|
| def __len__(self): |
| return len(self.indices) |
|
|
| def __getitem__(self, item): |
| index = self.indices[item] |
| timestamp = date.fromisoformat(str(self.data["timestamps"][index])) |
| observed = self.data["observed_anomaly"].copy() |
| precision = self.data["precision"].copy() |
| rng = np.random.default_rng(self.seed + index) |
| original_valid = np.isfinite(observed[index]) & self.data["ocean_mask"] |
| withheld = original_valid & (rng.random(original_valid.shape) < self.cloud_fraction) |
| observed[index, withheld] = np.nan |
| precision[index, withheld] = 0.0 |
| inputs = build_input(observed, precision, index, self.data["longitude"], |
| self.data["latitude"], timestamp.timetuple().tm_yday) |
| target = self.data["sst_anomaly"][index][None].astype(np.float32) |
| return torch.from_numpy(inputs), torch.from_numpy(target), torch.from_numpy(withheld[None]) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| rank = int(os.environ.get("RANK", "0")) |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| distributed = world_size > 1 |
| if distributed: |
| backend = "nccl" if torch.cuda.is_available() else "gloo" |
| torch.distributed.init_process_group(backend=backend) |
| requested = config["runtime"]["device"] |
| use_cuda = torch.cuda.is_available() and requested != "cpu" |
| device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu") |
| if use_cuda: |
| torch.cuda.set_device(device) |
| torch.manual_seed(int(config["seed"]) + rank) |
| path = ROOT / config["data"]["root"] / config["data"]["file"] |
| with np.load(path) as loaded: |
| arrays = {key: loaded[key] for key in loaded.files} |
| train_count = min(int(config["data"]["train_samples"]), len(arrays["timestamps"]) - 1) |
| dataset = SSTDataset(arrays, range(train_count), config["data"]["random_cloud_fraction"], config["seed"]) |
| sampler = DistributedSampler(dataset, shuffle=True, seed=int(config["seed"])) if distributed else None |
| loader = DataLoader(dataset, batch_size=int(config["training"]["batch_size"]), |
| shuffle=sampler is None, sampler=sampler, |
| num_workers=int(config["runtime"]["num_workers"])) |
| model = DINCAE(**config["model"]).to(device) |
| if distributed: |
| model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None) |
| optimizer = torch.optim.Adam(model.parameters(), lr=float(config["training"]["learning_rate"])) |
| history = [] |
| for epoch in range(int(config["training"]["epochs"])): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| model.train() |
| total, batches = 0.0, 0 |
| for inputs, target, mask in loader: |
| inputs, target, mask = inputs.to(device), target.to(device), mask.to(device) |
| optimizer.zero_grad(set_to_none=True) |
| output = model(inputs) |
| loss = masked_gaussian_nll(output, target, mask, |
| config["training"]["gamma"], config["training"]["delta"]) |
| loss.backward() |
| optimizer.step() |
| total, batches = total + float(loss.detach()), batches + 1 |
| history.append({"epoch": epoch + 1, "masked_gaussian_nll": total / max(batches, 1)}) |
| if rank == 0: |
| print(f"epoch={epoch + 1} loss={history[-1]['masked_gaussian_nll']:.6f}") |
| if rank == 0: |
| raw_model = model.module if distributed else model |
| checkpoint = ROOT / config["paths"]["checkpoint"] |
| checkpoint.parent.mkdir(parents=True, exist_ok=True) |
| torch.save({"model": raw_model.state_dict(), "model_config": config["model"], |
| "gamma": config["training"]["gamma"], "delta": config["training"]["delta"], |
| "format_version": "dincae_checkpoint_v1"}, checkpoint) |
| metrics = ROOT / config["paths"]["training_metrics"] |
| metrics.parent.mkdir(parents=True, exist_ok=True) |
| metrics.write_text(json.dumps({"history": history, "world_size": world_size, |
| "parameter_count": sum(p.numel() for p in raw_model.parameters())}, indent=2)) |
| print(f"checkpoint={checkpoint}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
