"""Train reduced TerraMind dual-scale masked token prediction.""" 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.terramind import TerraMind class TerraMindDataset(Dataset): def __init__(self, path, config, training=False): self.data = np.load(path) self.config = config self.training = training if str(self.data["format_version"]) != config["data"]["format_version"]: raise ValueError("incompatible data format") source = int(config["data"]["source_size"]) for name, channels in config["data"]["pixel_modalities"].items(): if self.data[f"pixel_{name}"].shape[1:] != (int(channels), source, source): raise ValueError(f"{name} does not preserve the TerraMesh source dimensions") def __len__(self): return len(self.data["pixel_s2l2a"]) def __getitem__(self, index): source, model_size = int(self.config["data"]["source_size"]), int(self.config["data"]["model_size"]) top = np.random.randint(0, source - model_size + 1) if self.training else (source - model_size) // 2 left = np.random.randint(0, source - model_size + 1) if self.training else (source - model_size) // 2 item = {} for name in self.config["data"]["pixel_modalities"]: values = self.data[f"pixel_{name}"][index, :, top:top + model_size, left:left + model_size] item[f"pixel_{name}"] = torch.from_numpy(values).float() patch = int(self.config["data"]["patch_size"]) vocab = int(self.config["model"]["engineering_vocab_size"]) image_tokens = { "s2l2a": item["pixel_s2l2a"].mean(dim=0), "s1grd": item["pixel_s1grd"].mean(dim=0), "s1rtc": item["pixel_s1rtc"].mean(dim=0), "dem": item["pixel_dem"].mean(dim=0), "ndvi": (item["pixel_s2l2a"][7] - item["pixel_s2l2a"][3]) / (item["pixel_s2l2a"][7] + item["pixel_s2l2a"][3]).abs().clamp_min(1e-3), "lulc": torch.from_numpy(self.data["token_map_lulc"][index, top:top + model_size, left:left + model_size]).float(), } for name, values in image_tokens.items(): pooled = torch.nn.functional.avg_pool2d(values[None, None], patch, stride=patch)[0, 0] minimum, maximum = pooled.amin(), pooled.amax() item[f"token_{name}"] = torch.round((pooled - minimum) / (maximum - minimum).clamp_min(1e-6) * (vocab - 1)).long().flatten() item["token_coords"] = torch.from_numpy(self.data["coords"][index]).long() item["token_caption"] = torch.from_numpy(self.data["caption"][index]).long() return item def device_from_config(config, rank=0): if config["runtime"]["device"] == "auto": return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu") return torch.device(config["runtime"]["device"]) def unpack(batch, config, device): pixels = {name: batch[f"pixel_{name}"].to(device) for name in config["data"]["pixel_modalities"]} tokens = {name: batch[f"token_{name}"].to(device) for name in config["data"]["token_modalities"]} return pixels, tokens def sample_task(pixels, tokens, config): pixel_count = random.randint(int(config["model"]["min_pixel_modalities"]), min(int(config["model"]["max_pixel_modalities"]), len(pixels))) selected_pixels = random.sample(list(pixels), pixel_count) target_count = random.randint(1, max(1, len(tokens) // 2)) targets = random.sample(list(tokens), target_count) candidates = [name for name in tokens if name not in targets] token_count = random.randint(int(config["model"]["min_token_modalities"]), min(int(config["model"]["max_token_modalities"]), len(candidates))) inputs = random.sample(candidates, token_count) return {name: pixels[name] for name in selected_pixels}, inputs, targets def main(): config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) torch.manual_seed(int(config["seed"])) random.seed(int(config["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) dataset = TerraMindDataset(ROOT / config["data"]["root"] / "train.npz", config, training=True) sampler = DistributedSampler(dataset, shuffle=True) if distributed else None loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler, shuffle=sampler is None, num_workers=int(config["train"]["num_workers"])) model = TerraMind(config["data"]["pixel_modalities"], config["data"]["token_modalities"], config["model"]).to(device) if distributed: model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) optimizer = torch.optim.AdamW(model.parameters(), lr=float(config["train"]["learning_rate"]), weight_decay=float(config["train"]["weight_decay"])) history = [] for epoch in range(int(config["train"]["epochs"])): total, steps = 0.0, 0 model.train() for batch in loader: pixels, tokens = unpack(batch, config, device) selected_pixels, input_tokens, targets = sample_task(pixels, tokens, config) output = model(selected_pixels, tokens, targets, input_tokens, apply_input_mask=True) optimizer.zero_grad(set_to_none=True) output["loss"].backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total += float(output["loss"].detach()) steps += 1 metrics = {"epoch": epoch + 1, "cross_modal_token_ce": total / max(steps, 1)} history.append(metrics) if rank == 0: print(f"epoch={epoch + 1} cross_modal_token_ce={metrics['cross_modal_token_ce']:.6f}") if rank == 0: checkpoint, metrics_path = ROOT / config["paths"]["checkpoint"], ROOT / config["paths"]["training_metrics"] checkpoint.parent.mkdir(parents=True, exist_ok=True) metrics_path.parent.mkdir(parents=True, exist_ok=True) state = model.module.state_dict() if distributed else model.state_dict() torch.save({"model": state, "model_config": config["model"], "pixel_modalities": config["data"]["pixel_modalities"], "token_modalities": config["data"]["token_modalities"], "format_version": config["data"]["format_version"]}, checkpoint) metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n") if distributed: torch.distributed.destroy_process_group() if __name__ == "__main__": main()