Fixes issue with model reconstruction of the upper half of the image & saves model checkpoint in s3

environment.yml

@@ -18,6 +18,7 @@ dependencies:
   - pytorch~=2.1.0 *cuda12*  # [linux]
   - python~=3.11.0
   - rioxarray~=0.15.0
+  - s3fs~=2024.3.1
   - scikit-image~=0.22.0
   - scikit-learn~=1.4.0
   - stackstac~=0.5.0

src/model_clay.py

@@ -57,6 +57,7 @@ def __init__(  # noqa: PLR0913
         mask_ratio,
         image_size,
         patch_size,
+        shuffle,
         dim,
         depth,
         heads,
@@ -73,6 +74,7 @@ def __init__(  # noqa: PLR0913
         self.mask_ratio = mask_ratio
         self.image_size = image_size
         self.patch_size = patch_size
+        self.shuffle = shuffle
         self.dim = dim
         self.band_groups = band_groups
         self.num_spatial_patches = (image_size // patch_size) ** 2
@@ -243,7 +245,7 @@ def mask_out(self, patches):
             GL == self.num_patches
         ), f"Expected {self.num_patches} patches, got {GL} patches."
 
-        if self.training:  # Shuffle the patches
+        if self.shuffle:  # Shuffle the patches
             noise = torch.randn((B, GL), device=patches.device)  # [B GL]
         else:  # Don't shuffle useful for interpolation & inspection of embeddings
             noise = rearrange(
@@ -560,6 +562,7 @@ def __init__(  # noqa: PLR0913
         mask_ratio,
         image_size,
         patch_size,
+        shuffle,
         # ENCODER
         dim,
         depth,
@@ -590,12 +593,14 @@ def __init__(  # noqa: PLR0913
         self.mask_ratio = mask_ratio
         self.image_size = image_size
         self.patch_size = patch_size
+        self.shuffle = shuffle
         self.band_groups = band_groups
 
         self.encoder = Encoder(
             mask_ratio=mask_ratio,
             image_size=image_size,
             patch_size=patch_size,
+            shuffle=shuffle,
             dim=dim,
             depth=depth,
             heads=heads,
@@ -786,12 +791,20 @@ def __init__(  # noqa: PLR0913
         mask_ratio=0.75,
         image_size=512,
         patch_size=32,
+        shuffle=False,
         lr=1e-4,
         wd=0.05,
         b1=0.9,
         b2=0.95,
         embeddings_level: Literal["mean", "patch", "group"] = "mean",
-        band_groups=None,
+        band_groups={
+            "rgb": (2, 1, 0),
+            "rededge": (3, 4, 5, 7),
+            "nir": (6,),
+            "swir": (8, 9),
+            "sar": (10, 11),
+            "dem": (12,),
+        },
     ):
         super().__init__()
         self.save_hyperparameters(logger=True)
@@ -806,6 +819,7 @@ def __init__(  # noqa: PLR0913
                 "mask_ratio": mask_ratio,
                 "image_size": image_size,
                 "patch_size": patch_size,
+                "shuffle": shuffle,
             }
             if band_groups:
                 model_args["band_groups"] = band_groups
@@ -826,7 +840,7 @@ def configure_optimizers(self):
             betas=(self.hparams.b1, self.hparams.b2),
         )
         scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
-            optimizer, T_0=1000, T_mult=2, eta_min=self.hparams.lr * 10, last_epoch=-1
+            optimizer, T_0=1000, T_mult=2, eta_min=self.hparams.lr * 100, last_epoch=-1
         )
 
         return {

trainer.py

@@ -36,25 +36,26 @@ def cli_main(
         "devices": "auto",
         "strategy": "ddp",
         "precision": "bf16-mixed",
-        "log_every_n_steps": 1,
+        "log_every_n_steps": 50,
         "max_epochs": 100,
         "accumulate_grad_batches": 5,
+        "default_root_dir": "s3://clay-model-ckpt/v0.2/",
         "callbacks": [
             ModelCheckpoint(
-                dirpath="checkpoints/",
+                dirpath="s3://clay-model-ckpt/v0.2/",
                 auto_insert_metric_name=False,
-                filename="mae_epoch-{epoch:02d}_val-loss-{val/loss:.2f}",
+                filename="mae_epoch-{epoch:02d}_val-loss-{val/loss:.4f}",
                 monitor="val/loss",
                 mode="min",
                 save_last=True,
                 save_top_k=2,
-                save_weights_only=True,
+                save_weights_only=False,
                 verbose=True,
             ),
             LearningRateMonitor(logging_interval="step"),
             LogIntermediatePredictions(),
         ],
-        "logger": [WandbLogger(project="CLAY-v0", log_model=False)],
+        "logger": [WandbLogger(entity="devseed-gaia", project="clay", log_model=False)],
         "plugins": [AsyncCheckpointIO()],
     },
     args: ArgsType = None,