"""Train or fine-tune ClimODE with OneScience ERA5 data.""" from __future__ import annotations import argparse import json import os import random import sys from pathlib import Path import numpy as np import torch import torch.distributed as dist import torch.nn as nn import yaml from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler # Allow ``python scripts/train.py`` to resolve project-local packages. PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from model.climode import ClimODE, load_checkpoint from scripts.data_loader import ClimODESeriesDataset, load_constants from scripts.velocity import fit_velocity_cache, load_velocity_cache def set_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def _device(value: str | None) -> torch.device: if value: return torch.device(value) return torch.device("cuda" if torch.cuda.is_available() else "cpu") def _init_distributed(backend: str) -> tuple[bool, int, int]: world_size = int(os.environ.get("WORLD_SIZE", "1")) if world_size == 1: return False, 0, 1 if not dist.is_initialized(): dist.init_process_group(backend=backend) return True, dist.get_rank(), world_size def _nll(mean: torch.Tensor, std: torch.Tensor, truth: torch.Tensor, var_coeff: float) -> torch.Tensor: distribution = torch.distributions.Normal(mean, 1.0e-3 + std) return (-distribution.log_prob(truth)).mean() + var_coeff * (std.square()).sum() def _load_yaml(path: Path) -> dict: with path.open("r", encoding="utf-8") as handle: return yaml.safe_load(handle) def _resolve(path: str | Path) -> Path: value = Path(path) return value if value.is_absolute() else PROJECT_ROOT / value def _parse_years(value: str | None, fallback: list[int]) -> list[int]: if value is None: return list(fallback) years = [int(item.strip()) for item in value.split(",") if item.strip()] if not years: raise ValueError("year override must contain at least one integer") return years def _model_from_args(config: dict, args: argparse.Namespace, device: torch.device) -> nn.Module: model_cfg = config["model"] use_pretrained = bool(getattr(args, "use_pretrained", False)) pretrained_checkpoint = getattr(args, "pretrained_checkpoint", None) if args.mode == "resume": checkpoint = args.checkpoint or _resolve(model_cfg["default_checkpoint"]) model = load_checkpoint(checkpoint, map_location="cpu") elif args.mode == "finetune" or use_pretrained: checkpoint = args.checkpoint if checkpoint is None and (use_pretrained or pretrained_checkpoint is not None): checkpoint = pretrained_checkpoint or model_cfg.get("pretrained_checkpoint") if checkpoint is None: raise ValueError( "finetune requires --checkpoint, or explicitly pass " "--use-pretrained [--pretrained-checkpoint PATH]" ) checkpoint = _resolve(checkpoint) if not checkpoint.is_file(): raise FileNotFoundError(f"Checkpoint not found: {checkpoint}") model = load_checkpoint(checkpoint, map_location="cpu") else: model = ClimODE( num_channels=5, const_channels=2, out_types=5, method=args.solver or model_cfg.get("solver", "euler"), use_attention=model_cfg.get("use_attention", True), use_uncertainty=model_cfg.get("use_uncertainty", True), use_positional_encoder=model_cfg.get("use_positional_encoder", False), ) return model.to(device) def _run_epoch( model, loader, velocity, constants, lat, lon, device, optimizer, var_coeff, max_batches, atol, rtol, ): training = optimizer is not None model.train(training) total = 0.0 count = 0 for batch_index, batch in enumerate(loader): if max_batches is not None and batch_index >= max_batches: break observations = batch["observations"].squeeze(0).to(device) time_steps = batch["time_steps"].squeeze(0).to(device) sequence_index = int(batch["sequence_index"].item()) past_velocity = velocity[sequence_index].to(device) target = observations initial = observations[0].unsqueeze(1) model_core = model.module if isinstance(model, DistributedDataParallel) else model model_core.update_param([past_velocity, constants, lat, lon]) if training: optimizer.zero_grad(set_to_none=True) with torch.set_grad_enabled(training): mean, std, _ = model(time_steps, initial, atol=atol, rtol=rtol) loss = _nll(mean, std, target, var_coeff) loss = loss + 0.001 * sum(parameter.square().sum() for parameter in model.parameters()) if training: loss.backward() optimizer.step() total += float(loss.detach()) count += 1 if dist.is_initialized(): totals = torch.tensor([total, float(count)], dtype=torch.float64, device=device) dist.all_reduce(totals, op=dist.ReduceOp.SUM) total, count = float(totals[0].item()), int(totals[1].item()) return total / max(count, 1), count def _prepare_velocity( dataset, constants: torch.Tensor, lat: torch.Tensor, lon: torch.Tensor, path: Path, epochs: int, learning_rate: float, smoothing_alpha: float, kernel_sigma: float, distributed: bool, rank: int, ) -> torch.Tensor: """Build a split cache once, then let every DDP rank read the same result.""" if path.is_file(): return load_velocity_cache(path, len(dataset)) if distributed: if rank == 0: fit_velocity_cache( dataset, constants, lat, lon, path, epochs=epochs, learning_rate=learning_rate, smoothing_alpha=smoothing_alpha, kernel_sigma=kernel_sigma, ) dist.barrier() return load_velocity_cache(path, len(dataset)) return fit_velocity_cache( dataset, constants, lat, lon, path, epochs=epochs, learning_rate=learning_rate, smoothing_alpha=smoothing_alpha, kernel_sigma=kernel_sigma, ) def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml" ) parser.add_argument("--mode", choices=["scratch", "finetune", "resume"], default=None) parser.add_argument("--checkpoint", type=Path, default=None) parser.add_argument( "--use-pretrained", action="store_true", help="Explicitly initialize from the official pretrained checkpoint", ) parser.add_argument( "--pretrained-checkpoint", type=Path, default=None, help="Override model.pretrained_checkpoint when --use-pretrained is set", ) parser.add_argument("--solver", choices=["euler", "rk4", "dopri5", "dopri8", "midpoint"], default=None) parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--sequence-length", type=int, default=None) parser.add_argument("--velocity-epochs", type=int, default=None) parser.add_argument("--velocity-cache", type=Path, default=None) parser.add_argument("--data-dir", type=Path, default=None, help="Override data.data_dir") parser.add_argument("--stats-dir", type=Path, default=None, help="Override data.stats_dir") parser.add_argument("--static-file", type=Path, default=None, help="Override data.static_file") parser.add_argument("--checkpoint-dir", type=Path, default=None) parser.add_argument("--log-file", type=Path, default=None) parser.add_argument("--device", type=str, default=None) parser.add_argument("--max-batches", type=int, default=None) parser.add_argument("--seed", type=int, default=None) parser.add_argument("--train-years", type=str, default=None, help="Comma-separated year override") parser.add_argument("--val-years", type=str, default=None, help="Comma-separated year override") args = parser.parse_args() args.config = _resolve(args.config) args.checkpoint = _resolve(args.checkpoint) if args.checkpoint is not None else None args.pretrained_checkpoint = ( _resolve(args.pretrained_checkpoint) if args.pretrained_checkpoint is not None else None ) args.data_dir = _resolve(args.data_dir) if args.data_dir is not None else None args.stats_dir = _resolve(args.stats_dir) if args.stats_dir is not None else None args.static_file = _resolve(args.static_file) if args.static_file is not None else None args.velocity_cache = ( _resolve(args.velocity_cache) if args.velocity_cache is not None else None ) args.checkpoint_dir = ( _resolve(args.checkpoint_dir) if args.checkpoint_dir is not None else None ) args.log_file = _resolve(args.log_file) if args.log_file is not None else None config = _load_yaml(args.config) model_cfg, data_cfg, vel_cfg, train_cfg = config["model"], config["data"], config["velocity"], config["training"] args.mode = args.mode or train_cfg.get("mode", "scratch") if args.mode == "resume" and args.use_pretrained: raise ValueError("--use-pretrained cannot be combined with --mode resume") if ( args.pretrained_checkpoint is not None and args.mode not in {"finetune"} and not args.use_pretrained ): raise ValueError( "--pretrained-checkpoint requires --use-pretrained or " "--mode finetune" ) if args.mode == "scratch" and args.checkpoint is not None: raise ValueError( "--checkpoint is ignored in scratch mode; use --mode resume or " "--mode finetune explicitly" ) if args.use_pretrained and args.mode == "scratch": args.mode = "finetune" args.solver = args.solver or model_cfg.get("solver", "euler") args.sequence_length = args.sequence_length or data_cfg.get("sequence_length", 8) args.velocity_epochs = args.velocity_epochs if args.velocity_epochs is not None else vel_cfg.get("epochs", 200) args.max_batches = args.max_batches if args.max_batches is not None else train_cfg.get("max_batches") set_seed(args.seed if args.seed is not None else train_cfg.get("seed", 42)) distributed, rank, world_size = _init_distributed(train_cfg.get("ddp_backend", "nccl")) device = _device(args.device) if distributed and device.type == "cuda": device = torch.device("cuda", int(os.environ.get("LOCAL_RANK", "0"))) if device.type == "cuda": if device.index is None: device = torch.device("cuda", 0) torch.cuda.set_device(device) root = _resolve(args.data_dir or data_cfg["data_dir"]) stats_dir = _resolve(args.stats_dir or data_cfg.get("stats_dir", root / "static")) train_set = ClimODESeriesDataset( root, _parse_years(args.train_years, data_cfg["train_years"]), stats_dir=stats_dir, model_size=(data_cfg["model_height"], data_cfg["model_width"]), sequence_length=args.sequence_length, normalize=data_cfg.get("normalize", True), ) val_set = ClimODESeriesDataset( root, _parse_years(args.val_years, data_cfg["val_years"]), stats_dir=stats_dir, model_size=(data_cfg["model_height"], data_cfg["model_width"]), sequence_length=args.sequence_length, normalize=data_cfg.get("normalize", True), ) train_sampler = DistributedSampler(train_set, shuffle=True) if distributed else None val_sampler = DistributedSampler(val_set, shuffle=False) if distributed else None train_loader = DataLoader(train_set, batch_size=1, sampler=train_sampler, shuffle=train_sampler is None, num_workers=data_cfg["dataloader"]["num_workers"]) val_loader = DataLoader(val_set, batch_size=1, sampler=val_sampler, shuffle=False, num_workers=data_cfg["dataloader"]["num_workers"]) static_file = _resolve(args.static_file or data_cfg["static_file"]) constants, lat, lon = load_constants(static_file, (data_cfg["model_height"], data_cfg["model_width"])) constants, lat, lon = constants.to(device), lat.unsqueeze(0).to(device), lon.unsqueeze(0).to(device) # Relative paths follow the project working directory, matching the # config/checkpoint conventions used by the reference earth projects. velocity_root = _resolve(vel_cfg["cache_dir"]) velocity_path = args.velocity_cache or (velocity_root / "train.pt") train_velocity = _prepare_velocity( train_set, constants, lat.squeeze(0).cpu(), lon.squeeze(0).cpu(), velocity_path, epochs=args.velocity_epochs, learning_rate=vel_cfg["learning_rate"], smoothing_alpha=vel_cfg["smoothing_alpha"], kernel_sigma=vel_cfg["kernel_sigma"], distributed=distributed, rank=rank, ) val_velocity_path = velocity_path.with_name("val.pt") val_velocity = _prepare_velocity( val_set, constants, lat.squeeze(0).cpu(), lon.squeeze(0).cpu(), val_velocity_path, epochs=args.velocity_epochs, learning_rate=vel_cfg["learning_rate"], smoothing_alpha=vel_cfg["smoothing_alpha"], kernel_sigma=vel_cfg["kernel_sigma"], distributed=distributed, rank=rank, ) model = _model_from_args(config, args, device) if distributed: model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None) lr = train_cfg.get("finetune_learning_rate", 5.0e-5) if args.mode == "finetune" else model_cfg.get("learning_rate", 5.0e-4) optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=model_cfg.get("weight_decay", 1.0e-5)) epochs = args.epochs or (train_cfg.get("finetune_epochs", 40) if args.mode == "finetune" else train_cfg.get("epochs", 300)) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs) start_epoch = 0 if args.mode == "resume": resume_path = args.checkpoint or _resolve(model_cfg["default_checkpoint"]) try: resume = torch.load(resume_path, map_location="cpu", weights_only=True) except TypeError: resume = torch.load(resume_path, map_location="cpu") state = resume.get("model", resume.get("state_dict")) (model.module if isinstance(model, DistributedDataParallel) else model).load_state_dict(state) if "optimizer" in resume: optimizer.load_state_dict(resume["optimizer"]) if "scheduler" in resume: scheduler.load_state_dict(resume["scheduler"]) start_epoch = int(resume.get("epoch", -1)) + 1 checkpoint_dir = _resolve(args.checkpoint_dir or model_cfg["checkpoint_dir"]) checkpoint_dir.mkdir(parents=True, exist_ok=True) log_path = _resolve(args.log_file or train_cfg.get("log_file", "./result/train.jsonl")) log_path.parent.mkdir(parents=True, exist_ok=True) best_val = float("inf") for epoch in range(start_epoch, epochs): if train_sampler is not None: train_sampler.set_epoch(epoch) var_coeff = 1.0e-3 if epoch == 0 else 2.0 * scheduler.get_last_lr()[0] train_loss, train_count = _run_epoch( model, train_loader, train_velocity, constants, lat, lon, device, optimizer, var_coeff, args.max_batches, model_cfg["atol"], model_cfg["rtol"] ) with torch.no_grad(): val_loss, val_count = _run_epoch( model, val_loader, val_velocity, constants, lat, lon, device, None, var_coeff, args.max_batches, model_cfg["atol"], model_cfg["rtol"] ) scheduler.step() record = {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "train_batches": train_count, "val_batches": val_count, "lr": scheduler.get_last_lr()[0]} if rank == 0: with log_path.open("a", encoding="utf-8") as handle: handle.write(json.dumps(record) + "\n") if val_loss < best_val: best_val = val_loss torch.save({"model": (model.module if isinstance(model, DistributedDataParallel) else model).state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "epoch": epoch}, checkpoint_dir / "model_bak.pth") print(json.dumps(record)) if distributed: dist.destroy_process_group() if __name__ == "__main__": main()