Update SkySense model package

conf/config.yaml

@@ -1,30 +1,38 @@
 seed: 42
 data:
   root: data
-  protocol: synthetic_skysense_npz
-  train_samples: 24
-  test_samples: 8
-  image_size: 32
+  protocol: skysense_multires_multitemporal_v2
+  train_samples: 2
+  test_samples: 1
+  hr_size: 224
+  s2_size: 64
+  s1_size: 64
   num_classes: 6
   hr_channels: 3
   s2_channels: 10
   s1_channels: 2
-  s2_timesteps: 4
-  s1_timesteps: 3
+  hr_timesteps: 1
+  s2_timesteps: 20
+  s1_timesteps: 10
 model:
-  embed_dim: 32
-  patch_size: 4
-  temporal_depth: 2
+  embed_dim: 24
+  hr_patch_size: 16
+  s2_patch_size: 8
+  s1_patch_size: 8
+  temporal_depth: 1
   temporal_heads: 4
   num_regions: 16
   prototypes_per_region: 4
 train:
-  epochs: 3
-  batch_size: 4
+  epochs: 1
+  batch_size: 1
   learning_rate: 0.001
   weight_decay: 0.0001
   alignment_weight: 0.05
   num_workers: 0
+  amp: true
+runtime:
+  device: cpu
 paths:
   checkpoint: result/checkpoints/skysense.pt
   training_metrics: result/training/metrics.json

config.json

@@ -5,40 +5,60 @@
     "SkySense"
   ],
   "framework": "PyTorch",
-  "domain": "earth-science",
-  "task": "remote-sensing-multimodal-semantic-segmentation",
+  "domain": "earth-observation",
+  "task": "remote-sensing-semantic-segmentation",
   "implementation": {
     "entry_point": "model/skysense.py",
-    "scope": "compact multi-modal temporal SkySense reproduction"
+    "scope": "compact multimodal spatial-temporal encoder with modality fusion and regional geo-context prototypes"
   },
   "architecture": {
-    "family": "multi-modal temporal remote sensing encoder",
-    "input_format": "NPZ multimodal temporal arrays",
+    "family": "factorized multimodal spatial-temporal encoder with geo-context prototypes",
     "modalities": [
       "high-resolution RGB",
-      "Sentinel-2 multispectral",
-      "Sentinel-1 SAR"
+      "Sentinel-2",
+      "Sentinel-1"
     ],
-    "image_size": 32,
-    "embed_dim": 32,
-    "patch_size": 4,
-    "temporal_depth": 2,
+    "hr_shape": [
+      1,
+      3,
+      224,
+      224
+    ],
+    "sentinel_2_shape": [
+      20,
+      10,
+      64,
+      64
+    ],
+    "sentinel_1_shape": [
+      10,
+      2,
+      64,
+      64
+    ],
+    "embed_dim": 24,
+    "hr_patch_size": 16,
+    "sentinel_2_patch_size": 8,
+    "sentinel_1_patch_size": 8,
+    "temporal_depth": 1,
     "temporal_heads": 4,
-    "task": "semantic segmentation"
+    "num_regions": 16,
+    "prototypes_per_region": 4,
+    "num_classes": 6
   },
   "data": {
-    "dataset": "SkySense pretraining data",
-    "official_paper": "https://arxiv.org/abs/2312.10115",
-    "protocol": "synthetic_skysense_npz",
-    "synthetic_train_samples": 24,
-    "synthetic_test_samples": 8
+    "datasets": [
+      "High-resolution RGB imagery",
+      "Sentinel-1",
+      "Sentinel-2"
+    ],
+    "protocol": "skysense_multires_multitemporal_v2",
+    "input_format": "multimodal_npz",
+    "train_samples": 2,
+    "test_samples": 1
   },
-  "metrics": [
-    "pixel_accuracy",
-    "mean_iou",
-    "per_class_iou"
-  ],
   "configuration_sources": [
+    "configuration.json",
     "conf/config.yaml",
     "model/skysense.py"
   ]

configuration.json

@@ -1,9 +1,10 @@
 {
   "framework": "PyTorch",
+  "license": "other",
   "task": "remote_sensing_semantic_segmentation",
   "model": "SkySense",
   "input_format": "multimodal_npz",
-  "protocol": "synthetic_skysense_npz",
+  "protocol": "skysense_multires_multitemporal_v2",
   "default_config": "conf/config.yaml",
   "train": "scripts/train.py",
   "inference": "scripts/inference.py",

model/skysense.py

@@ -31,7 +31,9 @@ class SkySense(nn.Module):
         s2_channels=10,
         s1_channels=2,
         embed_dim=32,
-        patch_size=4,
+        hr_patch_size=16,
+        s2_patch_size=8,
+        s1_patch_size=8,
         temporal_depth=2,
         temporal_heads=4,
         num_regions=16,
@@ -39,11 +41,10 @@ class SkySense(nn.Module):
         num_classes=6,
     ):
         super().__init__()
-        self.patch_size = patch_size
         self.num_regions = num_regions
-        self.hr_encoder = SpatialEncoder(hr_channels, embed_dim, patch_size)
-        self.s2_encoder = SpatialEncoder(s2_channels, embed_dim, patch_size)
-        self.s1_encoder = SpatialEncoder(s1_channels, embed_dim, patch_size)
+        self.hr_encoder = SpatialEncoder(hr_channels, embed_dim, hr_patch_size)
+        self.s2_encoder = SpatialEncoder(s2_channels, embed_dim, s2_patch_size)
+        self.s1_encoder = SpatialEncoder(s1_channels, embed_dim, s1_patch_size)
         self.date_embedding = nn.Embedding(366, embed_dim)
         self.modality_embedding = nn.Parameter(torch.zeros(3, embed_dim))
         self.fusion_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
@@ -57,6 +58,12 @@ class SkySense(nn.Module):
             norm_first=True,
         )
         self.temporal_fusion = nn.TransformerEncoder(layer, temporal_depth)
+        modality_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim, nhead=temporal_heads, dim_feedforward=embed_dim * 4,
+            dropout=0.0, activation="gelu", batch_first=True, norm_first=True,
+        )
+        self.modality_fusion = nn.TransformerEncoder(modality_layer, 1)
+        self.modality_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
         self.prototypes = nn.Parameter(
             torch.randn(num_regions, prototypes_per_region, embed_dim) * 0.02
         )
@@ -68,9 +75,16 @@ class SkySense(nn.Module):
         nn.init.normal_(self.date_embedding.weight, std=0.02)
         nn.init.normal_(self.modality_embedding, std=0.02)
         nn.init.normal_(self.fusion_token, std=0.02)
+        nn.init.normal_(self.modality_token, std=0.02)
 
     def _add_context(self, features, dates, modality_index):
-        date_context = self.date_embedding(dates.clamp(0, 365)).unsqueeze(-1).unsqueeze(-1)
+        if dates.dtype != torch.long:
+            raise TypeError(f"dates must use torch.int64, got {dates.dtype}")
+        if dates.shape != features.shape[:2]:
+            raise ValueError(f"dates shape {tuple(dates.shape)} does not match image batch/time {tuple(features.shape[:2])}")
+        if torch.any((dates < 0) | (dates > 364)):
+            raise ValueError("dates must contain day-of-year values in [0, 364]")
+        date_context = self.date_embedding(dates).unsqueeze(-1).unsqueeze(-1)
         modality = self.modality_embedding[modality_index].view(1, 1, -1, 1, 1)
         return features + date_context + modality
 
@@ -81,16 +95,34 @@ class SkySense(nn.Module):
             self._add_context(self.s1_encoder(s1), dates_s1, 2),
         )
 
+    def _aggregate_time(self, features):
+        batch, time, dim, height, width = features.shape
+        sequence = features.permute(0, 3, 4, 1, 2).reshape(-1, time, dim)
+        token = self.fusion_token.expand(sequence.shape[0], -1, -1)
+        fused = self.temporal_fusion(torch.cat([token, sequence], dim=1))[:, 0]
+        return fused.reshape(batch, height, width, dim).permute(0, 3, 1, 2)
+
     def forward(self, hr, s2, s1, dates_hr, dates_s2, dates_s1, region):
+        if region.dtype != torch.long:
+            raise TypeError(f"region must use torch.int64, got {region.dtype}")
+        if region.shape != (hr.shape[0],):
+            raise ValueError(f"region must have shape [{hr.shape[0]}], got {tuple(region.shape)}")
+        if torch.any((region < 0) | (region >= self.num_regions)):
+            raise ValueError(f"region IDs must be in [0, {self.num_regions - 1}]")
         modality_features = self.encode_modalities(hr, s2, s1, dates_hr, dates_s2, dates_s1)
-        temporal = torch.cat(modality_features, dim=1)
-        batch, time, dim, out_height, out_width = temporal.shape
-        temporal = temporal.permute(0, 3, 4, 1, 2).reshape(-1, time, dim)
-        token = self.fusion_token.expand(temporal.shape[0], -1, -1)
-        fused = self.temporal_fusion(torch.cat([token, temporal], dim=1))[:, 0]
+        aggregated = [self._aggregate_time(feature) for feature in modality_features]
+        target_size = aggregated[0].shape[-2:]
+        aligned = [aggregated[0]] + [
+            F.interpolate(feature, size=target_size, mode="bilinear", align_corners=False)
+            for feature in aggregated[1:]
+        ]
+        batch, dim, out_height, out_width = aligned[0].shape
+        modalities = torch.stack(aligned, dim=1).permute(0, 3, 4, 1, 2).reshape(-1, 3, dim)
+        token = self.modality_token.expand(modalities.shape[0], -1, -1)
+        fused = self.modality_fusion(torch.cat([token, modalities], dim=1))[:, 0]
         fused = fused.reshape(batch, out_height, out_width, dim)
 
-        regional_prototypes = self.prototypes[region.remainder(self.num_regions)]
+        regional_prototypes = self.prototypes[region]
         query = F.normalize(fused, dim=-1)
         keys = F.normalize(regional_prototypes, dim=-1)
         attention = torch.einsum("bhwd,bpd->bhwp", query, keys).softmax(dim=-1)

scripts/fake_data.py

@@ -18,15 +18,39 @@ def load_config():
 def make_split(path, count, config, seed):
     rng = np.random.default_rng(seed)
     data = config["data"]
-    size = data["image_size"]
-    hr = rng.normal(size=(count, 1, data["hr_channels"], size, size)).astype("float32")
-    s2 = rng.normal(size=(count, data["s2_timesteps"], data["s2_channels"], size, size)).astype("float32")
-    s1 = rng.normal(size=(count, data["s1_timesteps"], data["s1_channels"], size, size)).astype("float32")
-    dates_hr = rng.integers(0, 365, size=(count, 1), dtype=np.int64)
+    hr_size = data["hr_size"]
+    low_size = data["s2_size"]
+    if low_size != data["s1_size"]:
+        raise ValueError("The fake generator expects equal S1/S2 grid sizes")
+    y, x = np.mgrid[0:hr_size, 0:hr_size].astype(np.float32) / (hr_size - 1)
+    low_y, low_x = np.mgrid[0:low_size, 0:low_size].astype(np.float32) / (low_size - 1)
+    land = np.stack([x, y, np.sin(4 * np.pi * x) * np.cos(3 * np.pi * y)], axis=0)
+    hr = np.empty((count, data["hr_timesteps"], data["hr_channels"], hr_size, hr_size), np.float32)
+    s2 = np.empty((count, data["s2_timesteps"], data["s2_channels"], low_size, low_size), np.float32)
+    s1 = np.empty((count, data["s1_timesteps"], data["s1_channels"], low_size, low_size), np.float32)
+    for sample in range(count):
+        phase = rng.uniform(0, 2 * np.pi)
+        hr[sample, 0] = land + rng.normal(0, 0.04, land.shape)
+        for time in range(data["s2_timesteps"]):
+            seasonal = np.sin(2 * np.pi * time / data["s2_timesteps"] + phase)
+            s2[sample, time] = np.stack([
+                0.5 + 0.2 * np.sin((band + 1) * low_x + phase) + 0.1 * seasonal * low_y
+                for band in range(data["s2_channels"])
+            ]) + rng.normal(0, 0.025, (data["s2_channels"], low_size, low_size))
+        for time in range(data["s1_timesteps"]):
+            seasonal = np.cos(2 * np.pi * time / data["s1_timesteps"] + phase)
+            s1[sample, time, 0] = -0.6 + 0.3 * np.sin(4 * np.pi * low_x) * np.cos(3 * np.pi * low_y) + 0.1 * seasonal
+            s1[sample, time, 1] = -0.8 + 0.2 * low_y - 0.1 * seasonal
+        hr[sample] = np.clip(hr[sample], -1, 1)
+        s2[sample] = np.clip(s2[sample], 0, 1)
+    dates_hr = rng.integers(0, 365, size=(count, data["hr_timesteps"]), dtype=np.int64)
     dates_s2 = rng.integers(0, 365, size=(count, data["s2_timesteps"]), dtype=np.int64)
     dates_s1 = rng.integers(0, 365, size=(count, data["s1_timesteps"]), dtype=np.int64)
     region = rng.integers(0, config["model"]["num_regions"], size=count, dtype=np.int64)
-    signal = hr[:, 0, 0] + 0.35 * s2[:, :, 3].mean(axis=1) - 0.2 * s1[:, :, 0].mean(axis=1)
+    indices = np.linspace(0, low_size - 1, hr_size).round().astype(int)
+    s2_hr = s2[:, :, 3].mean(axis=1)[:, indices][:, :, indices]
+    s1_hr = s1[:, :, 0].mean(axis=1)[:, indices][:, :, indices]
+    signal = hr[:, 0, 0] + 0.35 * s2_hr - 0.2 * s1_hr
     bins = np.quantile(signal, np.linspace(0, 1, data["num_classes"] + 1)[1:-1])
     labels = np.digitize(signal, bins).astype(np.int64)
     np.savez_compressed(
@@ -41,6 +65,9 @@ def make_split(path, count, config, seed):
         labels=labels,
         data_source=np.asarray("synthetic"),
         protocol=np.asarray(config["data"]["protocol"]),
+        band_order_hr=np.asarray(["R", "G", "B"]),
+        band_order_s2=np.asarray(["B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B11", "B12"]),
+        band_order_s1=np.asarray(["VV", "VH"]),
     )
 
 

scripts/inference.py

@@ -1,6 +1,7 @@
 """Run SkySense inference and save arrays for evaluation."""
 
 import importlib.util
+import argparse
 from pathlib import Path
 
 import numpy as np
@@ -19,6 +20,9 @@ def load_model_class():
 
 
 def main():
+    parser = argparse.ArgumentParser(description="Run batched SkySense segmentation inference")
+    parser.add_argument("--batch-size", type=int)
+    args = parser.parse_args()
     with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
         config = yaml.safe_load(handle)
     checkpoint_path = ROOT / config["paths"]["checkpoint"]
@@ -27,7 +31,8 @@ def main():
             f"Missing checkpoint: {checkpoint_path.relative_to(ROOT)}. "
             "Run `python scripts/train.py` first."
         )
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    use_accelerator = torch.cuda.is_available() and config["runtime"].get("device", "auto") != "cpu"
+    device = torch.device("cuda" if use_accelerator else "cpu")
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
     SkySense = load_model_class()
     model = SkySense(
@@ -47,15 +52,49 @@ def main():
         )
     archive = np.load(test_path)
     keys = ["hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region"]
-    tensors = {key: torch.from_numpy(archive[key]).to(device) for key in keys}
+    arrays = {key: archive[key] for key in keys}
+    expected = {
+        "hr": (config["data"]["hr_timesteps"], config["data"]["hr_channels"], config["data"]["hr_size"], config["data"]["hr_size"]),
+        "s2": (config["data"]["s2_timesteps"], config["data"]["s2_channels"], config["data"]["s2_size"], config["data"]["s2_size"]),
+        "s1": (config["data"]["s1_timesteps"], config["data"]["s1_channels"], config["data"]["s1_size"], config["data"]["s1_size"]),
+        "dates_hr": (config["data"]["hr_timesteps"],),
+        "dates_s2": (config["data"]["s2_timesteps"],),
+        "dates_s1": (config["data"]["s1_timesteps"],),
+        "region": (),
+    }
+    sample_count = len(arrays["hr"])
+    for key, shape in expected.items():
+        if len(arrays[key]) != sample_count or tuple(arrays[key].shape[1:]) != shape:
+            raise ValueError(f"Invalid test {key} shape {arrays[key].shape}; expected [N,{','.join(map(str, shape))}]")
+    for key in ("hr", "s2", "s1"):
+        if not np.issubdtype(arrays[key].dtype, np.floating):
+            raise TypeError(f"{key} must use a floating dtype")
+    for key in ("dates_hr", "dates_s2", "dates_s1", "region"):
+        if arrays[key].dtype != np.int64:
+            raise TypeError(f"{key} must use int64")
+    if any(np.any((arrays[key] < 0) | (arrays[key] > 364)) for key in ("dates_hr", "dates_s2", "dates_s1")):
+        raise ValueError("Test dates must be in [0, 364]")
+    if np.any((arrays["region"] < 0) | (arrays["region"] >= config["model"]["num_regions"])):
+        raise ValueError("Test region IDs are out of range")
+    labels = archive["labels"]
+    if labels.dtype != np.int64 or labels.shape != (sample_count, config["data"]["hr_size"], config["data"]["hr_size"]):
+        raise ValueError("Test labels must be int64 [N,hr_size,hr_size]")
+    batch_size = args.batch_size or config["train"]["batch_size"]
+    predictions = []
+    all_probabilities = []
     with torch.inference_mode():
-        output = model(tensors["hr"], tensors["s2"], tensors["s1"], tensors["dates_hr"], tensors["dates_s2"], tensors["dates_s1"], tensors["region"])
-        probabilities = output["logits"].softmax(dim=1)
+        for start in range(0, len(arrays["hr"]), batch_size):
+            tensors = {key: torch.from_numpy(value[start:start + batch_size]).to(device)
+                       for key, value in arrays.items()}
+            output = model(tensors["hr"], tensors["s2"], tensors["s1"], tensors["dates_hr"], tensors["dates_s2"], tensors["dates_s1"], tensors["region"])
+            probabilities = output["logits"].softmax(dim=1).cpu().numpy()
+            all_probabilities.append(probabilities)
+            predictions.append(probabilities.argmax(axis=1))
     output_dir = ROOT / config["paths"]["inference_dir"]
     output_dir.mkdir(parents=True, exist_ok=True)
-    np.save(output_dir / "predictions.npy", probabilities.argmax(dim=1).cpu().numpy())
-    np.save(output_dir / "probabilities.npy", probabilities.cpu().numpy())
-    np.save(output_dir / "targets.npy", archive["labels"])
+    np.save(output_dir / "predictions.npy", np.concatenate(predictions))
+    np.save(output_dir / "probabilities.npy", np.concatenate(all_probabilities))
+    np.save(output_dir / "targets.npy", labels)
     data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown"
     protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown"
     np.savez(output_dir / "metadata.npz", data_source=data_source, protocol=protocol)

scripts/result.py

@@ -46,12 +46,11 @@ def main():
     output_dir.mkdir(parents=True, exist_ok=True)
     with (output_dir / "metrics.json").open("w", encoding="utf-8") as handle:
         json.dump(metrics, handle, indent=2)
-    figure, axes = plt.subplots(2, 2, figsize=(7, 7))
-    for index, axis in enumerate(axes.flat):
-        sample = index // 2
-        image = targets[sample] if index % 2 == 0 else predictions[sample]
+    figure, axes = plt.subplots(1, 2, figsize=(8, 4))
+    for index, axis in enumerate(axes):
+        image = targets[0] if index == 0 else predictions[0]
         axis.imshow(image, vmin=0, vmax=classes - 1, cmap="terrain")
-        axis.set_title(("Target" if index % 2 == 0 else "Prediction") + f" {sample}")
+        axis.set_title("Target" if index == 0 else "Prediction")
         axis.axis("off")
     figure.tight_layout()
     figure.savefig(output_dir / "comparison.png", dpi=120)

scripts/train.py

@@ -47,13 +47,14 @@ class NPZDataset(Dataset):
         return {key: torch.as_tensor(value[index]) for key, value in self.arrays.items()}
 
 
-def setup_device():
+def setup_device(config):
     world_size = int(os.environ.get("WORLD_SIZE", "1"))
     local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    use_accelerator = torch.cuda.is_available() and config["runtime"].get("device", "auto") != "cpu"
     if world_size > 1:
-        backend = "nccl" if torch.cuda.is_available() else "gloo"
+        backend = "nccl" if use_accelerator else "gloo"
         dist.init_process_group(backend=backend)
-    if torch.cuda.is_available():
+    if use_accelerator:
         torch.cuda.set_device(local_rank)
         return torch.device("cuda", local_rank), world_size, local_rank
     return torch.device("cpu"), world_size, local_rank
@@ -65,7 +66,8 @@ def main():
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
-    device, world_size, local_rank = setup_device()
+    device, world_size, local_rank = setup_device(config)
+    amp_enabled = bool(config["train"].get("amp", True) and device.type == "cuda")
     dataset_path = ROOT / config["data"]["root"] / "train.npz"
     if not dataset_path.exists():
         raise FileNotFoundError(
@@ -73,10 +75,38 @@ def main():
             "Run `python scripts/fake_data.py` for a synthetic connectivity test."
         )
     dataset = NPZDataset(dataset_path)
+    expected = {
+        "hr": (config["data"]["hr_timesteps"], config["data"]["hr_channels"], config["data"]["hr_size"], config["data"]["hr_size"]),
+        "s2": (config["data"]["s2_timesteps"], config["data"]["s2_channels"], config["data"]["s2_size"], config["data"]["s2_size"]),
+        "s1": (config["data"]["s1_timesteps"], config["data"]["s1_channels"], config["data"]["s1_size"], config["data"]["s1_size"]),
+        "labels": (config["data"]["hr_size"], config["data"]["hr_size"]),
+        "dates_hr": (config["data"]["hr_timesteps"],),
+        "dates_s2": (config["data"]["s2_timesteps"],),
+        "dates_s1": (config["data"]["s1_timesteps"],),
+        "region": (),
+    }
+    for key, shape in expected.items():
+        if tuple(dataset.arrays[key].shape[1:]) != shape:
+            raise ValueError(f"Expected {key} shaped [N,{','.join(map(str, shape))}], got {dataset.arrays[key].shape}")
+        if len(dataset.arrays[key]) != len(dataset):
+            raise ValueError(f"Array {key} has {len(dataset.arrays[key])} samples, expected {len(dataset)}")
+    for key in ("hr", "s2", "s1"):
+        if not np.issubdtype(dataset.arrays[key].dtype, np.floating):
+            raise TypeError(f"{key} must use a floating dtype, got {dataset.arrays[key].dtype}")
+    for key in ("dates_hr", "dates_s2", "dates_s1", "region", "labels"):
+        if dataset.arrays[key].dtype != np.int64:
+            raise TypeError(f"{key} must use int64, got {dataset.arrays[key].dtype}")
+    for key in ("dates_hr", "dates_s2", "dates_s1"):
+        if np.any((dataset.arrays[key] < 0) | (dataset.arrays[key] > 364)):
+            raise ValueError(f"{key} must contain day-of-year values in [0, 364]")
+    if np.any((dataset.arrays["region"] < 0) | (dataset.arrays["region"] >= config["model"]["num_regions"])):
+        raise ValueError(f"region IDs must be in [0, {config['model']['num_regions'] - 1}]")
+    if np.any((dataset.arrays["labels"] < 0) | (dataset.arrays["labels"] >= config["data"]["num_classes"])):
+        raise ValueError(f"labels must be in [0, {config['data']['num_classes'] - 1}]")
     if local_rank == 0:
         print(
             f"data_source={dataset.data_source} protocol={dataset.protocol} "
-            f"samples={len(dataset)} image_size={config['data']['image_size']}x{config['data']['image_size']}"
+            f"samples={len(dataset)} hr={config['data']['hr_size']} s2={config['data']['s2_size']} s1={config['data']['s1_size']}"
         )
     sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
     loader = DataLoader(
@@ -97,24 +127,36 @@ def main():
     if world_size > 1:
         model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
     optimizer = torch.optim.AdamW(model.parameters(), lr=config["train"]["learning_rate"], weight_decay=config["train"]["weight_decay"])
+    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
     final_loss = float("nan")
+    final_segmentation = float("nan")
+    final_alignment = float("nan")
     for epoch in range(config["train"]["epochs"]):
         if sampler is not None:
             sampler.set_epoch(epoch)
         model.train()
-        total = 0.0
+        totals = torch.zeros(4, dtype=torch.float64, device=device)
         for batch in loader:
             batch = {key: value.to(device) for key, value in batch.items()}
-            output = model(batch["hr"], batch["s2"], batch["s1"], batch["dates_hr"], batch["dates_s2"], batch["dates_s1"], batch["region"])
-            base_model = model.module if hasattr(model, "module") else model
-            segmentation = torch.nn.functional.cross_entropy(output["logits"], batch["labels"])
-            alignment = base_model.cross_modal_alignment_loss(output["features"])
-            loss = segmentation + config["train"]["alignment_weight"] * alignment
             optimizer.zero_grad(set_to_none=True)
-            loss.backward()
-            optimizer.step()
-            total += loss.item()
-        final_loss = total / len(loader)
+            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=amp_enabled):
+                output = model(batch["hr"], batch["s2"], batch["s1"], batch["dates_hr"], batch["dates_s2"], batch["dates_s1"], batch["region"])
+                base_model = model.module if hasattr(model, "module") else model
+                segmentation = torch.nn.functional.cross_entropy(output["logits"], batch["labels"])
+                alignment = base_model.cross_modal_alignment_loss(output["features"])
+                loss = segmentation + config["train"]["alignment_weight"] * alignment
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            batch_size = batch["hr"].shape[0]
+            totals += torch.tensor([loss.item() * batch_size, segmentation.item() * batch_size,
+                                    alignment.item() * batch_size, batch_size],
+                                   dtype=torch.float64, device=device)
+        if world_size > 1:
+            dist.all_reduce(totals, op=dist.ReduceOp.SUM)
+        final_loss = float((totals[0] / totals[3]).item())
+        final_segmentation = float((totals[1] / totals[3]).item())
+        final_alignment = float((totals[2] / totals[3]).item())
         if local_rank == 0:
             print(f"epoch={epoch + 1} loss={final_loss:.6f}")
     if local_rank == 0:
@@ -128,6 +170,10 @@ def main():
             json.dumps(
                 {
                     "final_loss": final_loss,
+                    "segmentation_loss": final_segmentation,
+                    "alignment_loss": final_alignment,
+                    "backward_completed": True,
+                    "amp_enabled": amp_enabled,
                     "epochs": config["train"]["epochs"],
                     "samples": len(dataset),
                     "data_source": dataset.data_source,