File size: 5,134 Bytes
43dce84 | 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 | """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()
|