From e25d46f23a1e72f7e7477b8813e15a1dddb0cc86 Mon Sep 17 00:00:00 2001 From: SRM Date: Mon, 11 Mar 2024 13:03:49 +0530 Subject: [PATCH 1/5] Update params for v0.2 model run - Lr -> 1e-5 to 1e-5 - Data -> Size: 256 x 256, patchsize: 16 - Log checkpoints to s3 - Save model params along with optimizer & epoch state --- src/model_clay.py | 2 +- trainer.py | 23 +++++++++++++++++------ 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/src/model_clay.py b/src/model_clay.py index 7b9dc134..80bb2ac6 100644 --- a/src/model_clay.py +++ b/src/model_clay.py @@ -826,7 +826,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 { diff --git a/trainer.py b/trainer.py index 92b15a9a..d32f9025 100644 --- a/trainer.py +++ b/trainer.py @@ -36,19 +36,20 @@ 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=None, 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"), @@ -63,8 +64,18 @@ def cli_main( Command-line inteface to run CLAYModule with ClayDataModule. """ cli = LightningCLI( - model_class=CLAYModule, - datamodule_class=ClayDataModule, + model_class=CLAYModule( + model_size="small", + mask_ratio=0.75, + image_size=256, + patch_size=16, + lr=1e-5, + ), + datamodule_class=ClayDataModule( + data_dir="data", + batch_size=10, + num_workers=8, + ), save_config_callback=save_config_callback, seed_everything_default=seed_everything_default, trainer_defaults=trainer_defaults, From aa3321bd25955c84caf05738b97338a5e61959e6 Mon Sep 17 00:00:00 2001 From: SRM Date: Wed, 20 Mar 2024 16:22:22 +0000 Subject: [PATCH 2/5] Log to devseed-gaia account of wandb & save checkpoints on s3. --- trainer.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/trainer.py b/trainer.py index d32f9025..89c44e72 100644 --- a/trainer.py +++ b/trainer.py @@ -42,7 +42,7 @@ def cli_main( "default_root_dir": "s3://clay-model-ckpt/v0.2/", "callbacks": [ ModelCheckpoint( - dirpath=None, + dirpath="s3://clay-model-ckpt/v0.2/", auto_insert_metric_name=False, filename="mae_epoch-{epoch:02d}_val-loss-{val/loss:.4f}", monitor="val/loss", @@ -55,7 +55,7 @@ def cli_main( 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, @@ -64,18 +64,8 @@ def cli_main( Command-line inteface to run CLAYModule with ClayDataModule. """ cli = LightningCLI( - model_class=CLAYModule( - model_size="small", - mask_ratio=0.75, - image_size=256, - patch_size=16, - lr=1e-5, - ), - datamodule_class=ClayDataModule( - data_dir="data", - batch_size=10, - num_workers=8, - ), + model_class=CLAYModule, + datamodule_class=ClayDataModule, save_config_callback=save_config_callback, seed_everything_default=seed_everything_default, trainer_defaults=trainer_defaults, From beee502394357f5a41a4e4df3279c39457232bb6 Mon Sep 17 00:00:00 2001 From: SRM Date: Wed, 20 Mar 2024 16:23:24 +0000 Subject: [PATCH 3/5] Fix issue with not shuffling during validation run. Use shuffle=True while training & validation. --- src/model_clay.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/src/model_clay.py b/src/model_clay.py index 80bb2ac6..41fce35a 100644 --- a/src/model_clay.py +++ b/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 From fcbc8717ccecbe34d3fe9fa065f3db758f24f601 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 26 Mar 2024 05:24:28 +0000 Subject: [PATCH 4/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/trainer.py b/trainer.py index 89c44e72..3ec29810 100644 --- a/trainer.py +++ b/trainer.py @@ -49,7 +49,7 @@ def cli_main( mode="min", save_last=True, save_top_k=2, - save_weights_only=False, + save_weights_only=False, verbose=True, ), LearningRateMonitor(logging_interval="step"), From bbfe00ea4f3396886230de6a1b669d63b391ad89 Mon Sep 17 00:00:00 2001 From: SRM Date: Fri, 19 Apr 2024 15:53:50 +0530 Subject: [PATCH 5/5] Add s3fs to environment.yml --- environment.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/environment.yml b/environment.yml index fb0e61eb..9fa7f593 100644 --- a/environment.yml +++ b/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