| """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() |
|
|
