Upload SatMAE model package

.gitattributes

@@ -1,35 +1,13 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,51 @@
+seed: 42
+data:
+  root: data
+  protocol: fmow_rgb_temporal
+  mode: temporal
+  image_size: 32
+  channels: 3
+  frames: 3
+  timestamp_fields: [year_offset_2002, month_zero_based, hour]
+  num_classes: 10
+  train_samples: 32
+  test_samples: 8
+  mask_ratio: 0.75
+model:
+  architecture: satmae_vit_large_patch16
+  runtime_profile: smoke
+  image_size: 32
+  patch_size: 4
+  in_channels: 3
+  frames: 3
+  embed_dim: 64
+  encoder_depth: 2
+  encoder_heads: 4
+  decoder_dim: 32
+  decoder_depth: 1
+  decoder_heads: 4
+  mode: temporal
+  spectral_groups: [[0, 1], [2]]
+  norm_pix_loss: false
+  same_mask: false
+  spatial_mask: false
+training:
+  epochs: 2
+  batch_size: 4
+  accum_iter: 1
+  base_learning_rate: 0.001
+  learning_rate: null
+  min_learning_rate: 0.0
+  warmup_epochs: 1
+  weight_decay: 0.05
+  num_workers: 0
+  resume: null
+  save_every: 1
+runtime:
+  device: auto
+  amp: true
+paths:
+  checkpoint: result/checkpoints/satmae.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,69 @@
+{
+  "model_name": "SatMAE",
+  "model_type": "satmae",
+  "architectures": [
+    "SatMAE"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "remote-sensing-representation-learning",
+  "implementation": {
+    "entry_point": "model/satmae.py",
+    "scope": "masked autoencoding for temporal or grouped multispectral satellite imagery",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "temporal and grouped multispectral masked autoencoder",
+    "runtime_profile": "smoke",
+    "mode": "temporal",
+    "image_size": 32,
+    "patch_size": 4,
+    "in_channels": 3,
+    "frames": 3,
+    "embed_dim": 64,
+    "encoder_depth": 2,
+    "encoder_heads": 4,
+    "decoder_dim": 32,
+    "decoder_depth": 1,
+    "decoder_heads": 4,
+    "mask_ratio": 0.75,
+    "spectral_groups": [[0, 1], [2]],
+    "norm_pix_loss": false,
+    "same_mask": false,
+    "spatial_mask": false
+  },
+  "data": {
+    "datasets": [
+      "fMoW RGB",
+      "fMoW-Sentinel",
+      "NAIP",
+      "EuroSAT",
+      "BigEarthNet",
+      "SpaceNet v1"
+    ],
+    "protocol": "fmow_rgb_temporal",
+    "format": "NPZ",
+    "input_key": "images",
+    "input_shape": ["N", 3, 3, 32, 32],
+    "timestamp_key": "timestamps",
+    "timestamp_shape": ["N", 3, 3],
+    "timestamp_fields": ["year_offset_2002", "month_zero_based", "hour"],
+    "label_key": "labels",
+    "label_shape": ["N"],
+    "num_classes": 10,
+    "default_train_file": "data/train.npz",
+    "default_test_file": "data/test.npz",
+    "required_metadata": ["source_protocol", "data_source"]
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/satmae.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,14 @@
+{
+  "framework": "PyTorch",
+  "task": "remote_sensing_representation_learning",
+  "model": "SatMAE",
+  "input_format": "BTCHW with BT3 timestamps, or BCHW",
+  "protocol": "SatMAE fMoW temporal and grouped multispectral MAE",
+  "implementation": "clean-room paper-aligned PyTorch implementation",
+  "upstream_reference_commit": "0b210aceb37a14bbbd897110db5b104b3271d818",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/satmae.py

@@ -0,0 +1,495 @@
+"""Paper-aligned SatMAE model components.
+
+This is an original implementation of the architecture described in SatMAE.
+The upstream repository was used only as a behavioral reference; no upstream
+source text is incorporated here.
+"""
+
+import math
+from functools import partial
+
+import torch
+from torch import nn
+
+
+def _sincos_1d(values, dim):
+    """Return a fixed sine-cosine embedding for arbitrary scalar positions."""
+    if dim <= 0:
+        return values.new_zeros((*values.shape, 0))
+    pairs = (dim + 1) // 2
+    omega = torch.arange(pairs, device=values.device, dtype=torch.float32)
+    omega = torch.exp(-math.log(10000.0) * omega / max(pairs - 1, 1))
+    phase = values.to(torch.float32).unsqueeze(-1) * omega
+    return torch.cat((phase.sin(), phase.cos()), dim=-1)[..., :dim]
+
+
+def _sincos_2d(grid_size, dim):
+    """Return a fixed row-major 2D sine-cosine position embedding."""
+    rows, cols = torch.meshgrid(
+        torch.arange(grid_size, dtype=torch.float32),
+        torch.arange(grid_size, dtype=torch.float32),
+        indexing="ij",
+    )
+    row_dim = dim // 2
+    return torch.cat(
+        (_sincos_1d(rows.reshape(-1), row_dim),
+         _sincos_1d(cols.reshape(-1), dim - row_dim)),
+        dim=-1,
+    )
+
+
+def _timestamp_embedding(timestamps, dim):
+    """Encode either scalar times or fMoW ``[year, month, hour]`` tuples."""
+    if timestamps.ndim == 2:
+        return _sincos_1d(timestamps, dim)
+    if timestamps.ndim != 3 or timestamps.shape[-1] != 3:
+        raise ValueError("timestamps must have shape [B, T] or [B, T, 3]")
+    field_dims = [dim // 3] * 3
+    for index in range(dim % 3):
+        field_dims[index] += 1
+    return torch.cat(
+        [_sincos_1d(timestamps[..., index], field_dim)
+         for index, field_dim in enumerate(field_dims)],
+        dim=-1,
+    )
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, image_size, patch_size, in_channels, embed_dim):
+        super().__init__()
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size) ** 2
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, images):
+        if images.shape[-2:] != (self.image_size, self.image_size):
+            raise ValueError(
+                f"expected {self.image_size}x{self.image_size} images, "
+                f"got {tuple(images.shape[-2:])}"
+            )
+        return self.proj(images).flatten(2).transpose(1, 2)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attention = nn.MultiheadAttention(
+            dim, num_heads, dropout=0.0, bias=True, batch_first=True
+        )
+        self.norm2 = norm_layer(dim)
+        hidden_dim = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, dim)
+        )
+
+    def forward(self, tokens):
+        normalized = self.norm1(tokens)
+        tokens = tokens + self.attention(
+            normalized, normalized, normalized, need_weights=False
+        )[0]
+        return tokens + self.mlp(self.norm2(tokens))
+
+
+class SatMAE(nn.Module):
+    """Masked autoencoder for temporal or grouped multispectral imagery.
+
+    Temporal inputs use shape ``[B, T, C, H, W]`` and optional timestamps
+    ``[B, T]``. Multispectral inputs use shape ``[B, C, H, W]``.
+    """
+
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=16,
+        in_channels=3,
+        frames=3,
+        embed_dim=1024,
+        encoder_depth=24,
+        encoder_heads=16,
+        decoder_dim=512,
+        decoder_depth=8,
+        decoder_heads=16,
+        mlp_ratio=4.0,
+        mode="temporal",
+        spectral_groups=None,
+        mask_ratio=0.75,
+        norm_pix_loss=False,
+        same_mask=False,
+        spatial_mask=False,
+        temporal_embed_dim=None,
+        decoder_temporal_embed_dim=None,
+        channel_embed_dim=None,
+        decoder_channel_embed_dim=None,
+        norm_layer=None,
+    ):
+        super().__init__()
+        if image_size % patch_size:
+            raise ValueError("image_size must be divisible by patch_size")
+        if not 0.0 <= mask_ratio < 1.0:
+            raise ValueError("mask_ratio must be in [0, 1)")
+        if mode not in {"temporal", "multispectral"}:
+            raise ValueError("mode must be temporal or multispectral")
+        if embed_dim % encoder_heads or decoder_dim % decoder_heads:
+            raise ValueError("embedding dimensions must be divisible by head counts")
+
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.frames = frames
+        self.embed_dim = embed_dim
+        self.decoder_dim = decoder_dim
+        self.mode = mode
+        self.mask_ratio = mask_ratio
+        self.norm_pix_loss = norm_pix_loss
+        self.same_mask = same_mask
+        self.spatial_mask = spatial_mask
+        self.grid_size = image_size // patch_size
+        self.num_patches = self.grid_size ** 2
+
+        if mode == "temporal":
+            self.spectral_groups = None
+            self.patch_embed = PatchEmbed(
+                image_size, patch_size, in_channels, embed_dim
+            )
+            self.token_groups = frames
+            semantic_dim = temporal_embed_dim
+            if semantic_dim is None:
+                semantic_dim = min(128, max(2, embed_dim // 4))
+            decoder_semantic_dim = decoder_temporal_embed_dim
+            if decoder_semantic_dim is None:
+                decoder_semantic_dim = min(64, max(2, decoder_dim // 4))
+            prediction_dims = [patch_size ** 2 * in_channels]
+        else:
+            groups = spectral_groups or [list(range(in_channels))]
+            flattened = [channel for group in groups for channel in group]
+            if sorted(flattened) != list(range(in_channels)):
+                raise ValueError("spectral_groups must partition all input channels")
+            self.spectral_groups = tuple(tuple(group) for group in groups)
+            self.patch_embed = nn.ModuleList(
+                PatchEmbed(image_size, patch_size, len(group), embed_dim)
+                for group in self.spectral_groups
+            )
+            self.token_groups = len(self.spectral_groups)
+            semantic_dim = channel_embed_dim
+            if semantic_dim is None:
+                semantic_dim = min(256, max(2, embed_dim // 4))
+            decoder_semantic_dim = decoder_channel_embed_dim
+            if decoder_semantic_dim is None:
+                decoder_semantic_dim = min(128, max(2, decoder_dim // 4))
+            prediction_dims = [patch_size ** 2 * len(g) for g in self.spectral_groups]
+
+        if not 0 < semantic_dim < embed_dim:
+            raise ValueError("encoder semantic embedding dimension is invalid")
+        if not 0 < decoder_semantic_dim < decoder_dim:
+            raise ValueError("decoder semantic embedding dimension is invalid")
+        self.semantic_dim = semantic_dim
+        self.decoder_semantic_dim = decoder_semantic_dim
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
+        self.register_buffer(
+            "spatial_pos_embed",
+            _sincos_2d(self.grid_size, embed_dim - semantic_dim),
+            persistent=True,
+        )
+        self.register_buffer(
+            "decoder_spatial_pos_embed",
+            _sincos_2d(self.grid_size, decoder_dim - decoder_semantic_dim),
+            persistent=True,
+        )
+        if mode == "multispectral":
+            group_ids = torch.arange(self.token_groups, dtype=torch.float32)
+            self.register_buffer(
+                "group_embed", _sincos_1d(group_ids, semantic_dim), persistent=True
+            )
+            self.register_buffer(
+                "decoder_group_embed",
+                _sincos_1d(group_ids, decoder_semantic_dim),
+                persistent=True,
+            )
+
+        self.blocks = nn.ModuleList(
+            TransformerBlock(embed_dim, encoder_heads, mlp_ratio, norm_layer)
+            for _ in range(encoder_depth)
+        )
+        self.norm = norm_layer(embed_dim)
+        self.decoder_embed = nn.Linear(embed_dim, decoder_dim)
+        self.decoder_blocks = nn.ModuleList(
+            TransformerBlock(decoder_dim, decoder_heads, mlp_ratio, norm_layer)
+            for _ in range(decoder_depth)
+        )
+        self.decoder_norm = norm_layer(decoder_dim)
+        self.decoder_pred = nn.ModuleList(
+            nn.Linear(decoder_dim, output_dim) for output_dim in prediction_dims
+        )
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        patch_embeds = (
+            [self.patch_embed]
+            if isinstance(self.patch_embed, PatchEmbed)
+            else self.patch_embed
+        )
+        for patch_embed in patch_embeds:
+            nn.init.xavier_uniform_(patch_embed.proj.weight.flatten(1))
+            if patch_embed.proj.bias is not None:
+                nn.init.zeros_(patch_embed.proj.bias)
+        nn.init.normal_(self.cls_token, std=0.02)
+        nn.init.normal_(self.mask_token, std=0.02)
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def patchify(self, images):
+        if images.ndim != 4:
+            raise ValueError("patchify expects [B, C, H, W]")
+        batch, channels, height, width = images.shape
+        patch = self.patch_size
+        if height != width or height != self.image_size:
+            raise ValueError(f"expected square images of size {self.image_size}")
+        patches = images.reshape(
+            batch, channels, height // patch, patch, width // patch, patch
+        )
+        patches = patches.permute(0, 2, 4, 1, 3, 5)
+        return patches.reshape(batch, self.num_patches, channels * patch ** 2)
+
+    def unpatchify(self, patches, channels=None):
+        channels = channels or self.in_channels
+        batch = patches.shape[0]
+        patch = self.patch_size
+        expected = channels * patch ** 2
+        if patches.shape[1:] != (self.num_patches, expected):
+            raise ValueError("patch tensor has incompatible shape")
+        images = patches.reshape(
+            batch, self.grid_size, self.grid_size, channels, patch, patch
+        )
+        images = images.permute(0, 3, 1, 4, 2, 5)
+        return images.reshape(batch, channels, self.image_size, self.image_size)
+
+    def _random_masking(self, tokens, mask_ratio, share_spatial_mask):
+        batch, length, dim = tokens.shape
+        if share_spatial_mask:
+            units = self.num_patches
+            len_keep_units = int(units * (1.0 - mask_ratio))
+            noise = torch.rand(batch, units, device=tokens.device)
+            spatial_order = noise.argsort(dim=1)
+            kept = [spatial_order[:, :len_keep_units] + g * units
+                    for g in range(self.token_groups)]
+            removed = [spatial_order[:, len_keep_units:] + g * units
+                       for g in range(self.token_groups)]
+            ids_shuffle = torch.cat(kept + removed, dim=1)
+            len_keep = len_keep_units * self.token_groups
+        else:
+            len_keep = int(length * (1.0 - mask_ratio))
+            ids_shuffle = torch.rand(batch, length, device=tokens.device).argsort(dim=1)
+        ids_restore = ids_shuffle.argsort(dim=1)
+        ids_keep = ids_shuffle[:, :len_keep]
+        visible = torch.gather(tokens, 1, ids_keep.unsqueeze(-1).expand(-1, -1, dim))
+        mask = torch.ones(batch, length, device=tokens.device)
+        mask[:, :len_keep] = 0
+        mask = torch.gather(mask, 1, ids_restore)
+        return visible, mask, ids_restore
+
+    def _temporal_tokens(self, images, timestamps):
+        if images.ndim != 5:
+            raise ValueError("temporal mode expects images shaped [B, T, C, H, W]")
+        batch, frames, channels, _, _ = images.shape
+        if frames != self.frames or channels != self.in_channels:
+            raise ValueError(
+                f"expected T={self.frames}, C={self.in_channels}; got T={frames}, C={channels}"
+            )
+        if timestamps is None:
+            timestamps = torch.arange(frames, device=images.device).expand(batch, -1)
+        if timestamps.shape[:2] != (batch, frames):
+            raise ValueError(
+                f"timestamps must start with shape {(batch, frames)}, "
+                f"got {tuple(timestamps.shape)}"
+            )
+        spatial = self.spatial_pos_embed.to(dtype=images.dtype)
+        time = _timestamp_embedding(timestamps, self.semantic_dim).to(dtype=images.dtype)
+        position = torch.cat(
+            (spatial.view(1, 1, self.num_patches, -1).expand(batch, frames, -1, -1),
+             time.unsqueeze(2).expand(-1, -1, self.num_patches, -1)),
+            dim=-1,
+        ).reshape(batch, frames * self.num_patches, self.embed_dim)
+        tokens = torch.stack(
+            [self.patch_embed(images[:, frame]) for frame in range(frames)], dim=1
+        ).reshape(batch, frames * self.num_patches, self.embed_dim)
+        return tokens + position, timestamps
+
+    def _multispectral_tokens(self, images):
+        if images.ndim != 4 or images.shape[1] != self.in_channels:
+            raise ValueError(
+                f"multispectral mode expects images shaped [B, {self.in_channels}, H, W]"
+            )
+        spatial = self.spatial_pos_embed.to(dtype=images.dtype)
+        group = self.group_embed.to(dtype=images.dtype)
+        positions = torch.cat(
+            (spatial.view(1, self.num_patches, -1).expand(self.token_groups, -1, -1),
+             group.view(self.token_groups, 1, -1).expand(-1, self.num_patches, -1)),
+            dim=-1,
+        ).reshape(1, self.token_groups * self.num_patches, self.embed_dim)
+        tokens = torch.cat(
+            [embed(images[:, channels])
+             for embed, channels in zip(self.patch_embed, self.spectral_groups)],
+            dim=1,
+        )
+        return tokens + positions
+
+    def forward_encoder(self, images, timestamps=None, mask_ratio=None):
+        ratio = self.mask_ratio if mask_ratio is None else mask_ratio
+        if not 0.0 <= ratio < 1.0:
+            raise ValueError("mask_ratio must be in [0, 1)")
+        if self.mode == "temporal":
+            tokens, timestamps = self._temporal_tokens(images, timestamps)
+            shared = self.same_mask
+        else:
+            tokens = self._multispectral_tokens(images)
+            shared = self.spatial_mask
+        tokens, mask, ids_restore = self._random_masking(tokens, ratio, shared)
+        cls = self.cls_token.expand(tokens.shape[0], -1, -1)
+        tokens = torch.cat((cls, tokens), dim=1)
+        for block in self.blocks:
+            tokens = block(tokens)
+        return self.norm(tokens), mask, ids_restore, timestamps
+
+    def _decoder_positions(self, batch, timestamps, dtype, device):
+        spatial = self.decoder_spatial_pos_embed.to(device=device, dtype=dtype)
+        if self.mode == "temporal":
+            semantic = _timestamp_embedding(timestamps, self.decoder_semantic_dim).to(dtype=dtype)
+        else:
+            semantic = self.decoder_group_embed.to(device=device, dtype=dtype)
+            semantic = semantic.unsqueeze(0).expand(batch, -1, -1)
+        position = torch.cat(
+            (spatial.view(1, 1, self.num_patches, -1).expand(batch, self.token_groups, -1, -1),
+             semantic.unsqueeze(2).expand(-1, -1, self.num_patches, -1)),
+            dim=-1,
+        )
+        return position.reshape(batch, self.token_groups * self.num_patches, self.decoder_dim)
+
+    def forward_decoder(self, latent, ids_restore, timestamps=None):
+        tokens = self.decoder_embed(latent)
+        mask_tokens = self.mask_token.expand(
+            tokens.shape[0], ids_restore.shape[1] + 1 - tokens.shape[1], -1
+        )
+        restored = torch.cat((tokens[:, 1:], mask_tokens), dim=1)
+        restored = torch.gather(
+            restored, 1, ids_restore.unsqueeze(-1).expand(-1, -1, self.decoder_dim)
+        )
+        positions = self._decoder_positions(
+            tokens.shape[0], timestamps, tokens.dtype, tokens.device
+        )
+        tokens = torch.cat((tokens[:, :1], restored + positions), dim=1)
+        for block in self.decoder_blocks:
+            tokens = block(tokens)
+        decoded = self.decoder_norm(tokens)[:, 1:]
+
+        if self.mode == "temporal":
+            return [self.decoder_pred[0](decoded)]
+        decoded = decoded.reshape(
+            decoded.shape[0], self.token_groups, self.num_patches, self.decoder_dim
+        )
+        return [head(decoded[:, index]) for index, head in enumerate(self.decoder_pred)]
+
+    def _targets(self, images):
+        if self.mode == "temporal":
+            return [torch.cat(
+                [self.patchify(images[:, frame]) for frame in range(self.frames)], dim=1
+            )]
+        return [self.patchify(images[:, group]) for group in self.spectral_groups]
+
+    def forward_loss(self, targets, predictions, mask):
+        losses = []
+        if self.mode == "temporal":
+            pairs = [(targets[0], predictions[0], mask)]
+        else:
+            group_mask = mask.reshape(mask.shape[0], self.token_groups, self.num_patches)
+            pairs = [
+                (target, prediction, group_mask[:, index])
+                for index, (target, prediction) in enumerate(zip(targets, predictions))
+            ]
+        removed = mask.new_zeros(())
+        total = mask.new_zeros(())
+        for target, prediction, patch_mask in pairs:
+            patch_loss = (prediction - target).square().mean(dim=-1)
+            total = total + (patch_loss * patch_mask).sum()
+            removed = removed + patch_mask.sum()
+            losses.append(patch_loss)
+        return total / removed.clamp_min(1), losses
+
+    def _normalize_targets(self, targets):
+        if not self.norm_pix_loss:
+            return targets
+        normalized = []
+        for target in targets:
+            mean = target.mean(dim=-1, keepdim=True)
+            variance = target.var(dim=-1, keepdim=True, unbiased=False)
+            normalized.append((target - mean) / torch.sqrt(variance + 1e-6))
+        return normalized
+
+    def _padded_outputs(self, tensors):
+        if self.mode == "temporal":
+            return tensors[0]
+        width = max(tensor.shape[-1] for tensor in tensors)
+        padded = []
+        for tensor in tensors:
+            if tensor.shape[-1] < width:
+                tensor = torch.nn.functional.pad(tensor, (0, width - tensor.shape[-1]))
+            padded.append(tensor)
+        return torch.cat(padded, dim=1)
+
+    def forward(self, images, timestamps=None, mask_ratio=None):
+        latent, mask, ids_restore, timestamps = self.forward_encoder(
+            images, timestamps, mask_ratio
+        )
+        predictions = self.forward_decoder(latent, ids_restore, timestamps)
+        targets = self._normalize_targets(self._targets(images))
+        loss, patch_losses = self.forward_loss(targets, predictions, mask)
+        return {
+            "loss": loss,
+            "prediction": self._padded_outputs(predictions),
+            "target": self._padded_outputs(targets),
+            "mask": mask.bool(),
+            "features": latent,
+            "ids_restore": ids_restore,
+            "group_predictions": predictions,
+            "group_targets": targets,
+            "patch_losses": patch_losses,
+        }
+
+
+def satmae_vit_base_patch16(**kwargs):
+    return SatMAE(
+        patch_size=16, embed_dim=768, encoder_depth=12, encoder_heads=12,
+        decoder_dim=512, decoder_depth=8, decoder_heads=16,
+        temporal_embed_dim=128, decoder_temporal_embed_dim=64,
+        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs
+    )
+
+
+def satmae_vit_large_patch16(**kwargs):
+    return SatMAE(
+        patch_size=16, embed_dim=1024, encoder_depth=24, encoder_heads=16,
+        decoder_dim=512, decoder_depth=8, decoder_heads=16,
+        temporal_embed_dim=128, decoder_temporal_embed_dim=64,
+        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs
+    )
+
+
+def satmae_vit_huge_patch14(**kwargs):
+    return SatMAE(
+        patch_size=14, embed_dim=1280, encoder_depth=32, encoder_heads=16,
+        decoder_dim=512, decoder_depth=8, decoder_heads=16,
+        temporal_embed_dim=128, decoder_temporal_embed_dim=64,
+        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs
+    )

scripts/fake_data.py

@@ -0,0 +1,45 @@
+"""Create temporary fMoW-style temporal tensors and labels."""
+import json
+from pathlib import Path
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    d = config["data"]
+    out = ROOT / d["root"]
+    out.mkdir(exist_ok=True)
+    rng = np.random.default_rng(config["seed"])
+    def make_split(samples):
+        shape = (samples, d["frames"], d["channels"], d["image_size"], d["image_size"])
+        images = rng.random(shape, dtype=np.float32)
+        timestamps = np.stack(
+            (
+                rng.integers(0, 21, size=(samples, d["frames"])),
+                rng.integers(0, 12, size=(samples, d["frames"])),
+                rng.integers(0, 24, size=(samples, d["frames"])),
+            ),
+            axis=-1,
+        ).astype(np.float32)
+        order = np.argsort(timestamps[..., 0] * 12 * 24 + timestamps[..., 1] * 24 + timestamps[..., 2], axis=1)
+        images = np.take_along_axis(images, order[:, :, None, None, None], axis=1)
+        timestamps = np.take_along_axis(timestamps, order[..., None], axis=1)
+        labels = rng.integers(d["num_classes"], size=samples, dtype=np.int64)
+        return images, timestamps, labels
+
+    train = make_split(d["train_samples"])
+    test = make_split(d["test_samples"])
+    np.savez_compressed(out / "train.npz", images=train[0], timestamps=train[1], labels=train[2])
+    np.savez_compressed(out / "test.npz", images=test[0], timestamps=test[1], labels=test[2])
+    (out / "format.json").write_text(json.dumps({
+        "format": "BTCHW",
+        "timestamp_format": "BT3: year_offset_2002, month_zero_based, hour",
+        "source_protocol": d["protocol"],
+        "data_source": "synthetic",
+    }, indent=2) + "\n")
+    print("created", out / "train.npz", out / "test.npz")
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,76 @@
+"""Run SatMAE masked reconstruction inference."""
+
+import argparse
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--data", type=Path, default=None)
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    parser.add_argument("--mask-ratio", type=float, default=None)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    config = yaml.safe_load(args.config.read_text())
+    spec = importlib.util.spec_from_file_location("satmae", ROOT / "model/satmae.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    model_args = {
+        key: value for key, value in config["model"].items()
+        if key not in {"architecture", "runtime_profile"}
+    }
+    model = module.SatMAE(**model_args)
+    checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model.load_state_dict(checkpoint["model"])
+    use_cuda = torch.cuda.is_available() and args.device != "cpu"
+    if args.device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but is unavailable")
+    device = torch.device("cuda" if use_cuda else "cpu")
+    model.to(device).eval()
+
+    data_path = args.data or ROOT / config["data"]["root"] / "test.npz"
+    archive = np.load(data_path)
+    images = torch.from_numpy(archive["images"]).to(device)
+    timestamps = None
+    if "timestamps" in archive:
+        timestamps = torch.from_numpy(archive["timestamps"]).to(device)
+    with torch.inference_mode():
+        output = model(images, timestamps=timestamps, mask_ratio=args.mask_ratio)
+    output_dir = args.output_dir or ROOT / config["paths"]["inference_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "target": output["target"].cpu().numpy(),
+        "prediction": output["prediction"].cpu().numpy(),
+        "mask": output["mask"].cpu().numpy(),
+        "labels": archive["labels"],
+    }
+    if timestamps is not None:
+        payload["timestamps"] = timestamps.cpu().numpy()
+    for index, (prediction, target) in enumerate(zip(
+        output["group_predictions"], output["group_targets"]
+    )):
+        payload[f"prediction_group_{index}"] = prediction.cpu().numpy()
+        payload[f"target_group_{index}"] = target.cpu().numpy()
+    np.savez_compressed(output_dir / "reconstruction.npz", **payload)
+    print("inference=", output_dir / "reconstruction.npz")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,127 @@
+"""Evaluate SatMAE masked reconstruction across time and channels."""
+import argparse
+import json
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np, yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--input", type=Path, default=None)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+def unpatchify(patches, image_size, patch_size, channels):
+    side = image_size // patch_size
+    image = patches.reshape(side, side, channels, patch_size, patch_size)
+    return image.transpose(2, 0, 3, 1, 4).reshape(channels, image_size, image_size)
+
+
+def display_image(image):
+    image = image[:3].transpose(1, 2, 0)
+    low, high = float(image.min()), float(image.max())
+    return np.clip((image - low) / max(high - low, 1e-8), 0.0, 1.0)
+
+
+def main():
+    args = parse_args()
+    cfg = yaml.safe_load(args.config.read_text())
+    source = args.input or ROOT / cfg["paths"]["inference_dir"] / "reconstruction.npz"
+    if not source.exists(): raise FileNotFoundError("Run inference before evaluation")
+    a = np.load(source)
+    masked = a["mask"].astype(bool)
+    out = args.output_dir or ROOT / cfg["paths"]["evaluation_dir"]; out.mkdir(parents=True, exist_ok=True)
+    if cfg["model"]["mode"] == "multispectral":
+        groups = cfg["model"]["spectral_groups"]
+        group_mask = masked.reshape(masked.shape[0], len(groups), -1)
+        group_mse, masked_group_mse = [], []
+        weighted_error = 0.0
+        weighted_count = 0
+        masked_error_sum = 0.0
+        masked_count = 0
+        for index, group in enumerate(groups):
+            target = a[f"target_group_{index}"]
+            prediction = a[f"prediction_group_{index}"]
+            squared = (prediction - target) ** 2
+            patch_error = squared.mean(axis=-1)
+            group_mse.append(float(squared.mean()))
+            selected = group_mask[:, index]
+            masked_group_mse.append(float(patch_error[selected].mean()))
+            weighted_error += float(squared.sum())
+            weighted_count += squared.size
+            masked_error_sum += float(patch_error[selected].sum())
+            masked_count += int(selected.sum())
+        result = {
+            "masked_mse": masked_error_sum / max(masked_count, 1),
+            "reconstruction_mse": weighted_error / max(weighted_count, 1),
+            "group_mse": group_mse,
+            "masked_group_mse": masked_group_mse,
+            "data_source": "synthetic",
+            "protocol": cfg["data"]["protocol"],
+        }
+        (out / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+        print(json.dumps(result, indent=2)); print("evaluation=", out)
+        return
+
+    squared_error = (a["prediction"] - a["target"]) ** 2
+    patch_error = squared_error.mean(axis=-1)
+    error = float(squared_error.mean())
+    masked_error = float(patch_error[masked].mean()) if masked.any() else error
+    result = {"masked_mse": masked_error, "reconstruction_mse": error, "data_source": "synthetic", "protocol": cfg["data"]["protocol"]}
+    size = cfg["model"]["image_size"]; patch = cfg["model"]["patch_size"]; channels = cfg["model"]["in_channels"]
+    patch_count = (size // patch) ** 2
+    target = a["target"][0, :patch_count]
+    prediction = a["prediction"][0, :patch_count]
+    patch_mask = a["mask"][0, :patch_count]
+    masked_target = target.copy(); masked_target[patch_mask] = 0.0
+    panels = [
+        ("Original", unpatchify(target, size, patch, channels)),
+        ("Masked input", unpatchify(masked_target, size, patch, channels)),
+        ("Reconstruction", unpatchify(prediction, size, patch, channels)),
+    ]
+    figure, axes = plt.subplots(1, 3, figsize=(10, 3.4))
+    for axis, (title, image) in zip(axes, panels):
+        axis.imshow(display_image(image)); axis.set_title(title); axis.axis("off")
+    figure.tight_layout(); figure.savefig(out / "temporal_frame_reconstruction.png", dpi=160, bbox_inches="tight"); plt.close(figure)
+
+    frames = cfg["model"]["frames"] if cfg["model"]["mode"] == "temporal" else 1
+    frame_mse, masked_frame_mse = [], []
+    channel_mse = np.zeros(channels, dtype=np.float64)
+    for frame in range(frames):
+        start, end = frame * patch_count, (frame + 1) * patch_count
+        frame_target = a["target"][:, start:end]
+        frame_prediction = a["prediction"][:, start:end]
+        mse = float(np.mean((frame_prediction - frame_target) ** 2))
+        frame_mse.append(mse)
+        frame_mask = masked[:, start:end]
+        frame_patch_error = patch_error[:, start:end]
+        masked_frame_mse.append(float(frame_patch_error[frame_mask].mean()) if frame_mask.any() else mse)
+        shaped_error = ((frame_prediction - frame_target) ** 2).reshape(-1, channels, patch * patch).mean(axis=(0, 2))
+        channel_mse += shaped_error
+    channel_mse /= frames
+
+    figure, axis = plt.subplots(figsize=(6.2, 3.8))
+    frame_index = np.arange(1, frames + 1)
+    axis.plot(frame_index, frame_mse, marker="o", linewidth=2, label="All patches")
+    axis.plot(frame_index, masked_frame_mse, marker="s", linewidth=2, label="Masked patches")
+    axis.set(xlabel="Time frame", ylabel="MSE", title="Temporal Reconstruction Error")
+    axis.legend(); axis.grid(alpha=0.25); figure.tight_layout(); figure.savefig(out / "temporal_reconstruction_error.png", dpi=160); plt.close(figure)
+
+    figure, axis = plt.subplots(figsize=(6.2, 3.8))
+    axis.bar(np.arange(channels), channel_mse, color="#287271")
+    axis.set_xticks(np.arange(channels), [f"C{i + 1}" for i in range(channels)])
+    axis.set(xlabel="Input channel", ylabel="MSE", title="Channel Reconstruction Error")
+    figure.tight_layout(); figure.savefig(out / "spectral_band_reconstruction.png", dpi=160); plt.close(figure)
+
+    result["frame_mse"] = frame_mse
+    result["masked_frame_mse"] = masked_frame_mse
+    result["channel_mse"] = channel_mse.tolist()
+    (out / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2)); print("evaluation=", out)
+
+if __name__ == "__main__": main()

scripts/train.py

@@ -0,0 +1,236 @@
+"""Pre-train SatMAE with masked reconstruction; supports torchrun."""
+
+import argparse
+import importlib.util
+import json
+import math
+import os
+import random
+from contextlib import nullcontext
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class NPZDataset(Dataset):
+    def __init__(self, path, mode):
+        archive = np.load(path)
+        self.images = archive["images"]
+        self.timestamps = archive["timestamps"] if "timestamps" in archive else None
+        if mode == "temporal" and self.timestamps is None:
+            raise ValueError("temporal datasets must contain timestamps")
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        images = torch.from_numpy(self.images[index])
+        if self.timestamps is None:
+            return images, torch.empty(0)
+        return images, torch.from_numpy(self.timestamps[index])
+
+
+def load_model_class():
+    spec = importlib.util.spec_from_file_location("satmae", ROOT / "model/satmae.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.SatMAE
+
+
+def model_config(config):
+    return {
+        key: value for key, value in config["model"].items()
+        if key not in {"architecture", "runtime_profile"}
+    }
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--data", type=Path, default=None)
+    parser.add_argument("--output", type=Path, default=None)
+    parser.add_argument("--resume", type=Path, default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
+    return parser.parse_args()
+
+
+def cosine_learning_rate(progress, config, peak_lr):
+    warmup = config["warmup_epochs"]
+    if warmup > 0 and progress < warmup:
+        return peak_lr * progress / warmup
+    span = max(config["epochs"] - warmup, 1)
+    phase = min(max((progress - warmup) / span, 0.0), 1.0)
+    return config["min_learning_rate"] + 0.5 * (
+        peak_lr - config["min_learning_rate"]
+    ) * (1.0 + math.cos(math.pi * phase))
+
+
+def main():
+    args = parse_args()
+    config = yaml.safe_load(args.config.read_text())
+    train_config = config["training"]
+    if args.epochs is not None:
+        train_config["epochs"] = args.epochs
+    if args.batch_size is not None:
+        train_config["batch_size"] = args.batch_size
+
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    rank = int(os.environ.get("RANK", "0"))
+    distributed = world_size > 1
+    requested_device = args.device or config["runtime"]["device"]
+    use_cuda = torch.cuda.is_available() and requested_device != "cpu"
+    if requested_device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but is unavailable")
+    if distributed:
+        dist.init_process_group("nccl" if use_cuda else "gloo")
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    if use_cuda:
+        torch.cuda.set_device(local_rank)
+
+    seed = config["seed"] + rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    data_path = args.data or ROOT / config["data"]["root"] / "train.npz"
+    if not data_path.exists():
+        raise FileNotFoundError(f"training data not found: {data_path}")
+    dataset = NPZDataset(data_path, config["model"]["mode"])
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(
+        dataset,
+        batch_size=train_config["batch_size"],
+        shuffle=sampler is None,
+        sampler=sampler,
+        num_workers=train_config["num_workers"],
+        pin_memory=use_cuda,
+        drop_last=False,
+    )
+
+    model = load_model_class()(**model_config(config)).to(device)
+    model_without_ddp = model
+    if distributed:
+        model = DistributedDataParallel(
+            model, device_ids=[local_rank] if use_cuda else None
+        )
+        model_without_ddp = model.module
+
+    effective_batch = (
+        train_config["batch_size"] * train_config["accum_iter"] * world_size
+    )
+    peak_lr = train_config["learning_rate"]
+    if peak_lr is None:
+        peak_lr = train_config["base_learning_rate"] * effective_batch / 256
+    decay, no_decay = [], []
+    for name, parameter in model_without_ddp.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        (no_decay if parameter.ndim == 1 or name.endswith("bias") else decay).append(parameter)
+    optimizer = torch.optim.AdamW(
+        [
+            {"params": decay, "weight_decay": train_config["weight_decay"]},
+            {"params": no_decay, "weight_decay": 0.0},
+        ],
+        lr=peak_lr,
+        betas=(0.9, 0.95),
+    )
+    amp_enabled = bool(config["runtime"].get("amp", True) and use_cuda)
+    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
+    start_epoch = 0
+    history = []
+    resume_path = args.resume
+    if resume_path is None and train_config.get("resume"):
+        resume_path = ROOT / train_config["resume"]
+    if resume_path is not None:
+        checkpoint = torch.load(resume_path, map_location="cpu", weights_only=False)
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        if checkpoint.get("scaler") is not None:
+            scaler.load_state_dict(checkpoint["scaler"])
+        start_epoch = checkpoint["epoch"] + 1
+        history = checkpoint.get("history", [])
+
+    checkpoint_path = args.output or ROOT / config["paths"]["checkpoint"]
+    metrics_path = ROOT / config["paths"]["training_metrics"]
+    optimizer.zero_grad(set_to_none=True)
+    for epoch in range(start_epoch, train_config["epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        total_loss = 0.0
+        steps = len(loader)
+        for step, (images, timestamps) in enumerate(loader):
+            progress = epoch + step / max(steps, 1)
+            learning_rate = cosine_learning_rate(progress, train_config, peak_lr)
+            for group in optimizer.param_groups:
+                group["lr"] = learning_rate
+            images = images.to(device, non_blocking=use_cuda)
+            timestamps = timestamps.to(device, non_blocking=use_cuda)
+            timestamps = timestamps if timestamps.numel() else None
+            autocast = partial(torch.amp.autocast, "cuda") if amp_enabled else nullcontext
+            with autocast():
+                output = model(images, timestamps=timestamps)
+                loss = output["loss"] / train_config["accum_iter"]
+            if not torch.isfinite(loss):
+                raise ValueError(f"non-finite loss at epoch {epoch}, step {step}")
+            scaler.scale(loss).backward()
+            update = (step + 1) % train_config["accum_iter"] == 0 or step + 1 == steps
+            if update:
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+            total_loss += output["loss"].detach().item()
+
+        epoch_loss = total_loss / max(steps, 1)
+        record = {
+            "epoch": epoch + 1,
+            "reconstruction_loss": epoch_loss,
+            "learning_rate": optimizer.param_groups[0]["lr"],
+        }
+        history.append(record)
+        if rank == 0:
+            print(
+                f"epoch={epoch + 1} reconstruction_loss={epoch_loss:.6f} "
+                f"lr={record['learning_rate']:.3e}"
+            )
+            if (epoch + 1) % train_config["save_every"] == 0 or epoch + 1 == train_config["epochs"]:
+                checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+                torch.save(
+                    {
+                        "model": model_without_ddp.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                        "scaler": scaler.state_dict() if amp_enabled else None,
+                        "epoch": epoch,
+                        "history": history,
+                        "config": config,
+                    },
+                    checkpoint_path,
+                )
+
+    if rank == 0:
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.write_text(json.dumps({
+            "history": history,
+            "protocol": config["data"]["protocol"],
+            "data_source": "synthetic" if "synthetic" in data_path.name or (data_path.parent / "format.json").exists() else "provided",
+            "effective_batch_size": effective_batch,
+            "peak_learning_rate": peak_lr,
+        }, indent=2) + "\n")
+        print("checkpoint=", checkpoint_path)
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()