| """Train the ConvLSTM classifier with optional distributed data parallelism.""" |
|
|
| import json |
| import os |
| import random |
| import sys |
| 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.firecubenet import FireCubeNet |
|
|
|
|
| class WildfireDataset(Dataset): |
| def __init__(self, path, config): |
| self.data = np.load(path) |
| expected = config["data"] |
| if str(self.data["format_version"]) != expected["format_version"]: |
| raise ValueError("incompatible wildfire data format") |
| expected_shape = (int(expected["sequence_days"]), int(expected["channels"]), |
| int(expected["patch_height"]), int(expected["patch_width"])) |
| if self.data["inputs"].ndim != 5 or self.data["inputs"].shape[1:] != expected_shape: |
| raise ValueError(f"inputs must have shape [B,{','.join(map(str, expected_shape))}]") |
| count = len(self.data["inputs"]) |
| if self.data["labels"].shape != (count, 1): |
| raise ValueError("labels must have shape [B,1]") |
| if self.data["coords"].shape != (count, 2) or self.data["timestamps_unix_s"].shape != (count,): |
| raise ValueError("coords/timestamps shape mismatch") |
| if not np.isfinite(self.data["inputs"]).all() or not np.isfinite(self.data["labels"]).all(): |
| raise ValueError("inputs and labels must be finite") |
| if not np.isin(self.data["labels"], (0, 1)).all(): |
| raise ValueError("labels must be binary") |
| cover_sum = self.data["inputs"][:, :, 15:25].sum(axis=2) |
| if not np.allclose(cover_sum, 1.0, atol=1e-5): |
| raise ValueError("land-cover fractions must sum to one") |
|
|
| def __len__(self): |
| return len(self.data["labels"]) |
|
|
| def __getitem__(self, index): |
| return (torch.from_numpy(self.data["inputs"][index]).float(), |
| torch.from_numpy(self.data["labels"][index]).float()) |
|
|
|
|
| def device_from_config(config, local_rank=0): |
| requested = config["runtime"]["device"] |
| if requested == "auto": |
| return torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu") |
| return torch.device(requested) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| seed = int(config["seed"]) |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
| distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| if distributed: |
| torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo") |
| rank = torch.distributed.get_rank() if distributed else 0 |
| device = device_from_config(config, local_rank) |
| if device.type == "cuda": |
| torch.cuda.set_device(device) |
| dataset = WildfireDataset(ROOT / config["data"]["root"] / "train.npz", config) |
| sampler = DistributedSampler(dataset, shuffle=True, seed=seed) if distributed else None |
| loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), |
| shuffle=sampler is None, sampler=sampler, |
| num_workers=int(config["train"]["num_workers"])) |
| channel_mean = dataset.data["inputs"].mean(axis=(0, 1, 3, 4)).astype(np.float32) |
| channel_std = dataset.data["inputs"].std(axis=(0, 1, 3, 4)).clip(1e-6).astype(np.float32) |
| model = FireCubeNet(**config["model"]).to(device) |
| wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model |
| optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"]), |
| weight_decay=float(config["train"]["weight_decay"])) |
| criterion = torch.nn.BCEWithLogitsLoss() |
| mean = torch.from_numpy(channel_mean).to(device).view(1, 1, -1, 1, 1) |
| std = torch.from_numpy(channel_std).to(device).view(1, 1, -1, 1, 1) |
| history = [] |
| for epoch in range(int(config["train"]["epochs"])): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| total, samples = 0.0, 0 |
| wrapped.train() |
| for inputs, labels in loader: |
| inputs, labels = inputs.to(device), labels.to(device) |
| logits = wrapped((inputs - mean) / std) |
| loss = criterion(logits, labels) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(wrapped.parameters(), float(config["train"]["gradient_clip_norm"])) |
| optimizer.step() |
| total += float(loss.detach()) * len(inputs) |
| samples += len(inputs) |
| loss_sum = torch.tensor([total, samples], dtype=torch.float64, device=device) |
| if distributed: |
| torch.distributed.all_reduce(loss_sum) |
| if rank == 0: |
| history.append({"epoch": epoch + 1, "bce_with_logits": float(loss_sum[0] / loss_sum[1])}) |
| if rank == 0: |
| checkpoint_path = ROOT / config["paths"]["checkpoint"] |
| metrics_path = ROOT / config["paths"]["training_metrics"] |
| checkpoint_path.parent.mkdir(parents=True, exist_ok=True) |
| metrics_path.parent.mkdir(parents=True, exist_ok=True) |
| bare_model = wrapped.module if distributed else wrapped |
| torch.save({ |
| "model_state_dict": bare_model.state_dict(), |
| "optimizer_state_dict": optimizer.state_dict(), |
| "model_config": config["model"], "epoch": int(config["train"]["epochs"]), |
| "channel_mean": channel_mean, "channel_std": channel_std, |
| "format_version": config["data"]["format_version"], "seed": seed, |
| }, checkpoint_path) |
| metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n") |
| print(f"checkpoint={checkpoint_path.relative_to(ROOT)} final_loss={history[-1]['bce_with_logits']:.6f}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
