File size: 4,747 Bytes
4c4d99c | 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 | """Train the reduced Prithvi-EO-2.0 temporal-location MAE."""
import json
import os
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.prithvi_eo import PrithviEO2
class PrithviDataset(Dataset):
def __init__(self, path, config):
self.data = np.load(path)
self.config = config
if str(self.data["format_version"]) != config["data"]["format_version"]:
raise ValueError("incompatible data format")
expected = (
int(config["data"]["channels"]), int(config["data"]["frames"]),
int(config["data"]["image_size"]), int(config["data"]["image_size"]),
)
if self.data["pixels"].shape[1:] != expected:
raise ValueError(f"pixels have shape {self.data['pixels'].shape[1:]}, expected {expected}")
self.mean = torch.tensor(config["data"]["mean"], dtype=torch.float32)[:, None, None, None]
self.std = torch.tensor(config["data"]["std"], dtype=torch.float32)[:, None, None, None]
def __len__(self):
return len(self.data["pixels"])
def __getitem__(self, index):
pixels = torch.from_numpy(self.data["pixels"][index]).float()
return {
"pixels": (pixels - self.mean) / self.std,
"temporal": torch.from_numpy(self.data["temporal_coords"][index]).float(),
"location": torch.from_numpy(self.data["location_coords"][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())
torch.manual_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)
if device.type == "cuda":
torch.cuda.set_device(device)
dataset = PrithviDataset(ROOT / config["data"]["root"] / "train.npz", config)
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 = PrithviEO2(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"]), betas=(0.9, 0.95))
history = []
for epoch in range(int(config["train"]["epochs"])):
if sampler:
sampler.set_epoch(epoch)
model.train()
total, steps = 0.0, 0
for batch in loader:
output = model(batch["pixels"].to(device), batch["temporal"].to(device), batch["location"].to(device))
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, "masked_patch_mse": total / max(steps, 1)}
history.append(metrics)
if rank == 0:
print(f"epoch={epoch + 1} masked_patch_mse={metrics['masked_patch_mse']:.6f}")
if rank == 0:
checkpoint = ROOT / config["paths"]["checkpoint"]
metrics_path = 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"],
"format_version": config["data"]["format_version"]}, checkpoint)
metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
print(f"checkpoint={checkpoint.relative_to(ROOT)}")
if distributed:
torch.distributed.destroy_process_group()
if __name__ == "__main__":
main()
|