From 48748f15cd6366a8ca5843e0826ade794cca944d Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Mon, 15 Jul 2024 12:27:14 +0530 Subject: [PATCH 01/83] Compile clay model encoder --- src/__init__.py | 0 src/export.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++++ src/model.py | 16 ++++++------ src/utils.py | 5 +++- 4 files changed, 80 insertions(+), 9 deletions(-) create mode 100644 src/__init__.py create mode 100644 src/export.py diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/export.py b/src/export.py new file mode 100644 index 00000000..2adccbb9 --- /dev/null +++ b/src/export.py @@ -0,0 +1,68 @@ +from pathlib import Path + +import torch +from torch.export import Dim + +from src.model import ClayMAEModule + +CHECKPOINT_PATH = "checkpoints/clay-v1-base.ckpt" +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +# device = torch.device("cpu") + + +def get_data(): + # Load data + cube = torch.randn(128, 3, 224, 224).to(device) + time = torch.randn(128, 4).to(device) + latlon = torch.randn(128, 4).to(device) + waves = torch.randn(3).to(device) + gsd = torch.randn(1).to(device) + return cube, time, latlon, waves, gsd + + +def load_model(): + module = ClayMAEModule.load_from_checkpoint(CHECKPOINT_PATH) + encoder = module.model.encoder # Get the encoder + encoder = encoder.to(device) # Move to device + return encoder + + +def main(): + # Load data + cube, time, latlon, waves, gsd = get_data() + + # Load model + encoder = load_model() + + # Define dynamic shapes for model export + batch_size = Dim("batch_size", min=2, max=128) # Define batch size range + channel_bands = Dim("channel_bands", min=1, max=12) # Define channel bands range + + dynamic_shapes = { + "cube": {0: batch_size, 1: channel_bands}, + "time": {0: batch_size}, + "latlon": {0: batch_size}, + "waves": {0: channel_bands}, + "gsd": {0: None}, + } + + # Export model + exp_compiled_encoder = torch.export.export( + mod=encoder, + args=(cube, time, latlon, waves, gsd), + dynamic_shapes=dynamic_shapes, + strict=False, + ) + + # tensortrt compiled model + # trt_encoder = torch_tensorrt.dynamo.compile( + # exp_compiled_encoder, [cube, time, latlon, waves, gsd] + # ) + + # Save model + Path("checkpoints/compiled").mkdir(parents=True, exist_ok=True) + torch.export.save(exp_compiled_encoder, "checkpoints/compiled/encoder.pt") + + +if __name__ == "__main__": + main() diff --git a/src/model.py b/src/model.py index ee211b97..1ecf508e 100644 --- a/src/model.py +++ b/src/model.py @@ -160,14 +160,14 @@ def mask_out(self, patches): masked_matrix, ) # [B L:(1 - mask_ratio) D], [(1-mask_ratio)], [mask_ratio], [B L] - def forward(self, datacube): - cube, time, latlon, gsd, waves = ( - datacube["pixels"], # [B C H W] - datacube["time"], # [B 2] - datacube["latlon"], # [B 2] - datacube["gsd"], # 1 - datacube["waves"], # [N] - ) # [B C H W] + def forward(self, cube, time, latlon, waves, gsd): + # cube, time, latlon, gsd, waves = ( + # datacube["pixels"], # [B C H W] + # datacube["time"], # [B 2] + # datacube["latlon"], # [B 2] + # datacube["gsd"], # 1 + # datacube["waves"], # [N] + # ) # [B C H W] B, C, H, W = cube.shape diff --git a/src/utils.py b/src/utils.py index 539a2acd..1e35731f 100644 --- a/src/utils.py +++ b/src/utils.py @@ -11,6 +11,7 @@ def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype=torch.float32): assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb" omega = torch.arange(dim // 4) / (dim // 4 - 1) omega = 1.0 / (temperature**omega) + omega = omega.to(y.device) y = y.flatten()[:, None] * omega[None, :] x = x.flatten()[:, None] * omega[None, :] @@ -24,8 +25,9 @@ def posemb_sincos_2d_with_gsd( y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb" - omega = torch.arange(dim // 4) / (dim // 4 - 1) + omega = torch.arange(dim // 4, device=gsd.device) / (dim // 4 - 1) omega = 1.0 / (temperature ** (2 * omega / dim)) * (gsd / 1.0) # Adjusted for g + omega = omega.to(y.device) y = y.flatten()[:, None] * omega[None, :] x = x.flatten()[:, None] * omega[None, :] @@ -41,6 +43,7 @@ def posemb_sincos_1d(pos, dim, temperature: int = 10000, dtype=torch.float32): omega = torch.arange(dim // 2) / (dim // 2 - 1) omega = 1.0 / (temperature**omega) + omega = omega.to(pos.device) scaled_pos = pos[:, None] * omega[None, :] pe = torch.cat((scaled_pos.sin(), scaled_pos.cos()), dim=1) From 97eb19a23a177031a02d9d73ba5db417d4de3d5f Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 18 Jul 2024 17:37:52 +0530 Subject: [PATCH 02/83] Add benchmark & test files for the compiled clay encoder --- src/benchmark_encoder.py | 80 +++++++++++++++++++++++++++++++++++++ src/export.py | 65 ++++++++++++++++-------------- src/model.py | 9 ++++- src/test_encoder.py | 86 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 208 insertions(+), 32 deletions(-) create mode 100644 src/benchmark_encoder.py create mode 100644 src/test_encoder.py diff --git a/src/benchmark_encoder.py b/src/benchmark_encoder.py new file mode 100644 index 00000000..08c2b0ac --- /dev/null +++ b/src/benchmark_encoder.py @@ -0,0 +1,80 @@ +import argparse +import time +import warnings + +import torch + +warnings.filterwarnings("ignore") + +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def get_data(): + """ + Generate random data tensors for model input. + """ + cube = torch.randn(128, 3, 256, 256).to(DEVICE) + timestep = torch.randn(128, 4).to(DEVICE) + latlon = torch.randn(128, 4).to(DEVICE) + waves = torch.randn(3).to(DEVICE) + gsd = torch.randn(1).to(DEVICE) + return cube, timestep, latlon, waves, gsd + + +def load_exported_model(eager=True): + """ + Load the exported model from a file. + + Args: + eager (bool): Flag to decide whether to use eager mode or compiled mode. + """ + print("Loading exported model") + ep = torch.export.load("checkpoints/compiled/encoder.pt") + if eager: + model = ep.module() + else: + model = torch.compile(ep.module(), backend="inductor") + return model + + +def benchmark_model(model): + """ + Benchmark the model by running inference on randomly generated data. + + Args: + model: The model to benchmark. + """ + print("Benchmarking model") + start = time.time() + for i in range(20): + cube, timestep, latlon, waves, gsd = get_data() + with torch.inference_mode(): + out = model(cube, timestep, latlon, waves, gsd) + print( + f"Iteration {i}: Output shapes - {out[0].shape}, {out[1].shape}, {out[2].shape}, {out[3].shape}" # noqa E501 + ) + print("Time taken for inference: ", time.time() - start) + + +def run(eager=True): + """ + Run the exported model and benchmark it. + + Args: + eager (bool): Flag to decide whether to use eager mode or compiled mode. + """ + print("Running model") + model = load_exported_model(eager=eager) + benchmark_model(model) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Run benchmark for the exported model." + ) + parser.add_argument( + "--eager", action="store_true", help="Use eager mode for running the model." + ) + args = parser.parse_args() + + run(args.eager) diff --git a/src/export.py b/src/export.py index 2adccbb9..70a65f1e 100644 --- a/src/export.py +++ b/src/export.py @@ -1,3 +1,4 @@ +import warnings from pathlib import Path import torch @@ -5,38 +6,47 @@ from src.model import ClayMAEModule +warnings.filterwarnings("ignore") + CHECKPOINT_PATH = "checkpoints/clay-v1-base.ckpt" -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -# device = torch.device("cpu") +DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") +CHIP_SIZE = 256 def get_data(): - # Load data - cube = torch.randn(128, 3, 224, 224).to(device) - time = torch.randn(128, 4).to(device) - latlon = torch.randn(128, 4).to(device) - waves = torch.randn(3).to(device) - gsd = torch.randn(1).to(device) - return cube, time, latlon, waves, gsd + """ + Generate random data tensors for model input. + """ + cube = torch.randn(128, 3, CHIP_SIZE, CHIP_SIZE).to(DEVICE) + timestep = torch.randn(128, 4).to(DEVICE) + latlon = torch.randn(128, 4).to(DEVICE) + waves = torch.randn(3).to(DEVICE) + gsd = torch.randn(1).to(DEVICE) + return cube, timestep, latlon, waves, gsd def load_model(): - module = ClayMAEModule.load_from_checkpoint(CHECKPOINT_PATH) - encoder = module.model.encoder # Get the encoder - encoder = encoder.to(device) # Move to device + """ + Load the model from a checkpoint and prepare it for evaluation. + """ + module = ClayMAEModule.load_from_checkpoint( + CHECKPOINT_PATH, shuffle=False, mask_ratio=0.0 + ) + encoder = module.model.encoder.eval() # Get the encoder in eval mode + encoder = encoder.to(DEVICE) # Move to the appropriate device return encoder -def main(): - # Load data - cube, time, latlon, waves, gsd = get_data() - - # Load model +def export_model(): + """ + Export the model with dynamic shapes for deployment. + """ + cube, timestep, latlon, waves, gsd = get_data() encoder = load_model() # Define dynamic shapes for model export - batch_size = Dim("batch_size", min=2, max=128) # Define batch size range - channel_bands = Dim("channel_bands", min=1, max=12) # Define channel bands range + batch_size = Dim("batch_size", min=32, max=1200) + channel_bands = Dim("channel_bands", min=1, max=10) dynamic_shapes = { "cube": {0: batch_size, 1: channel_bands}, @@ -47,22 +57,17 @@ def main(): } # Export model - exp_compiled_encoder = torch.export.export( + ep = torch.export.export( mod=encoder, - args=(cube, time, latlon, waves, gsd), + args=(cube, timestep, latlon, waves, gsd), dynamic_shapes=dynamic_shapes, - strict=False, + strict=True, ) - # tensortrt compiled model - # trt_encoder = torch_tensorrt.dynamo.compile( - # exp_compiled_encoder, [cube, time, latlon, waves, gsd] - # ) - - # Save model + # Save the exported model Path("checkpoints/compiled").mkdir(parents=True, exist_ok=True) - torch.export.save(exp_compiled_encoder, "checkpoints/compiled/encoder.pt") + torch.export.save(ep, "checkpoints/compiled/encoder.pt") if __name__ == "__main__": - main() + export_model() diff --git a/src/model.py b/src/model.py index 1ecf508e..900b18b3 100644 --- a/src/model.py +++ b/src/model.py @@ -39,6 +39,10 @@ def __init__( # noqa: PLR0913 self.dim = dim self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02) + # Required to compile & export the model + self.grid_size = 256 // 8 + self.num_patches = self.grid_size**2 + self.patch_embedding = DynamicEmbedding( wave_dim=128, num_latent_tokens=128, @@ -64,8 +68,9 @@ def add_encodings(self, patches, time, latlon, gsd): """Add position encoding to the patches""" B, L, D = patches.shape - grid_size = int(math.sqrt(L)) - self.num_patches = grid_size**2 + # grid_size = int(math.sqrt(L)) + # self.num_patches = grid_size**2 + grid_size = self.grid_size pos_encoding = ( posemb_sincos_2d_with_gsd( diff --git a/src/test_encoder.py b/src/test_encoder.py new file mode 100644 index 00000000..14839197 --- /dev/null +++ b/src/test_encoder.py @@ -0,0 +1,86 @@ +import torch + +from src.datamodule import ClayDataModule + +# Load the pre-trained Clay encoder model +clay_encoder = torch.export.load("checkpoints/compiled/encoder.pt").module() + + +def load_batch(): + # Initialize the data module with appropriate parameters + dm = ClayDataModule( + data_dir="/home/ubuntu/data", + size=256, + metadata_path="configs/metadata.yaml", + batch_size=1, + num_workers=1, + ) + + # Setup the data module for the 'fit' stage + dm.setup(stage="fit") + metadata = dm.metadata + + # Get the training data loader and create an iterator + trn_dl = dm.train_dataloader() + iter_dl = iter(trn_dl) + + return iter_dl, metadata + + +def prepare_data(sensor, metadata, device): + """ + Load data from the sensor and transfer it to the specified device. + + Args: + - sensor (dict): Sensor data containing 'pixels', 'time', 'latlon', and 'platform'. + - metadata (dict): Metadata information for different platforms. + - device (torch.device): The device to which the data should be transferred. + + Returns: + - tuple: Transferred cube, timestep, latlon, waves, and gsd tensors. + """ + cube = sensor["pixels"] + timestep = sensor["time"] + latlon = sensor["latlon"] + platform = sensor["platform"][0] + + # Get wavelengths and ground sampling distance (gsd) from metadata + waves = torch.tensor(list(metadata[platform].bands.wavelength.values())) + gsd = torch.tensor([metadata[platform].gsd]) + + # Transfer data to the specified device + cube, timestep, latlon, waves, gsd = map( + lambda x: x.to(device), (cube, timestep, latlon, waves, gsd) + ) + return cube, timestep, latlon, waves, gsd + + +def main(): + dl, metadata = load_batch() + + # Fetch samples from the data loader + l8_c2l1 = next(dl) + l8_c2l2 = next(dl) + linz = next(dl) + naip = next(dl) + s1 = next(dl) + s2 = next(dl) + + # Perform inference with the Clay encoder model + with torch.no_grad(): + for sensor in (l8_c2l1, l8_c2l2, linz, naip, s1, s2): + # Load data and transfer to GPU + batch = prepare_data(sensor, metadata, torch.device("cuda")) + + # Get patch embeddings from the encoder model + patch_embeddings, *_ = clay_encoder(*batch) + + # Extract the class (CLS) embedding + cls_embedding = patch_embeddings[:, 0, :] + + # Print the platform and the shape of the CLS embedding + print(sensor["platform"][0], cls_embedding.shape) + + +if __name__ == "__main__": + main() From eba1867f1c5d443d76dd1369764807174bea52bf Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 24 Jul 2024 20:03:40 +0530 Subject: [PATCH 03/83] Revert changes to Encoder, don't change the API --- src/model.py | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/src/model.py b/src/model.py index 900b18b3..ee211b97 100644 --- a/src/model.py +++ b/src/model.py @@ -39,10 +39,6 @@ def __init__( # noqa: PLR0913 self.dim = dim self.cls_token = nn.Parameter(torch.randn(1, 1, dim) * 0.02) - # Required to compile & export the model - self.grid_size = 256 // 8 - self.num_patches = self.grid_size**2 - self.patch_embedding = DynamicEmbedding( wave_dim=128, num_latent_tokens=128, @@ -68,9 +64,8 @@ def add_encodings(self, patches, time, latlon, gsd): """Add position encoding to the patches""" B, L, D = patches.shape - # grid_size = int(math.sqrt(L)) - # self.num_patches = grid_size**2 - grid_size = self.grid_size + grid_size = int(math.sqrt(L)) + self.num_patches = grid_size**2 pos_encoding = ( posemb_sincos_2d_with_gsd( @@ -165,14 +160,14 @@ def mask_out(self, patches): masked_matrix, ) # [B L:(1 - mask_ratio) D], [(1-mask_ratio)], [mask_ratio], [B L] - def forward(self, cube, time, latlon, waves, gsd): - # cube, time, latlon, gsd, waves = ( - # datacube["pixels"], # [B C H W] - # datacube["time"], # [B 2] - # datacube["latlon"], # [B 2] - # datacube["gsd"], # 1 - # datacube["waves"], # [N] - # ) # [B C H W] + def forward(self, datacube): + cube, time, latlon, gsd, waves = ( + datacube["pixels"], # [B C H W] + datacube["time"], # [B 2] + datacube["latlon"], # [B 2] + datacube["gsd"], # 1 + datacube["waves"], # [N] + ) # [B C H W] B, C, H, W = cube.shape From 73171ddeb7f10adf0f5cd51be31542aa479f6f1a Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 24 Jul 2024 22:48:49 +0530 Subject: [PATCH 04/83] Add embedder to load clay encoder & save in onnx/ep format --- finetune/embedder/factory.py | 303 +++++++++++++++++++++++++++++++++++ 1 file changed, 303 insertions(+) create mode 100644 finetune/embedder/factory.py diff --git a/finetune/embedder/factory.py b/finetune/embedder/factory.py new file mode 100644 index 00000000..bf3ee6e4 --- /dev/null +++ b/finetune/embedder/factory.py @@ -0,0 +1,303 @@ +"""Export the Clay model to ONNX and pytorch ExportedProgram format. + +This script exports the Clay model to ONNX and pytorch ExportedProgram format +for deployment. The model is exported with dynamic shapes for inference. + +How to use: + +```bash +python -m finetune.embedder.factory \ + --img_size 256 \ + --ckpt_path checkpoints/clay-v1-base.ckpt \ + --device cuda \ + --name clay-v1-encoder.onnx \ + --onnx +# exports Clay encoder to ONNX format that can handle chips of size 256x256 +# for different sensors like Sentinel-2, Landsat-8, NAIP, LINZ & Sentinel 1. +``` + +```bash +python -m finetune.embedder.factory \ + --img_size 224 \ + --ckpt_path checkpoints/clay-v1-base.ckpt \ + --device cuda \ + --name clay-v1-encoder.pt2 \ + --ep +# exports Clay encoder to pytorch ExportedProgram format that can handle chips +# of size 224x224 for different sensors like Sentinel-2, Landsat-8, NAIP, LINZ +# & Sentinel 1. +``` + +""" + +import argparse +import re +import warnings +from pathlib import Path + +import torch +from einops import repeat +from torch import nn +from torch.export import Dim + +from src.model import Encoder +from src.utils import posemb_sincos_2d_with_gsd + +warnings.filterwarnings("ignore", category=UserWarning) + + +class EmbeddingEncoder(Encoder): + """Clay Encoder without mask and shuffle.""" + + def __init__( # noqa: PLR0913 + self, + img_size, + patch_size, + dim, + depth, + heads, + dim_head, + mlp_ratio, + ): + super().__init__( + mask_ratio=0.0, + shuffle=False, + patch_size=patch_size, + dim=dim, + depth=depth, + heads=heads, + dim_head=dim_head, + mlp_ratio=mlp_ratio, + ) + self.img_size = img_size + + # Using fixed grid size for inference + self.grid_size = img_size // patch_size + self.num_patches = self.grid_size**2 + + def add_encodings(self, patches, time, latlon, gsd): + """Add position encoding to the patches""" + B, L, D = patches.shape + + grid_size = self.grid_size + + pos_encoding = ( + posemb_sincos_2d_with_gsd( + h=grid_size, + w=grid_size, + dim=(self.dim - 8), + gsd=gsd, + ) + .to(patches.device) + .detach() + ) # [L (D - 8)] + + time_latlon = torch.hstack((time, latlon)).to(patches.device).detach() # [B 8] + + pos_encoding = repeat(pos_encoding, "L D -> B L D", B=B) # [B L (D - 8)] + time_latlon = repeat(time_latlon, "B D -> B L D", L=L) # [B L 8] + pos_metadata_encoding = torch.cat( + (pos_encoding, time_latlon), dim=-1 + ) # [B L D] + + patches = patches + pos_metadata_encoding # [B L D] + [B L D] -> [B L D] + return patches # [B L D] + + # def forward(self, cube, time, latlon, waves, gsd): + def forward(self, datacube): + cube, time, latlon, gsd, waves = ( + datacube["pixels"], # [B C H W] + datacube["time"], # [B 2] + datacube["latlon"], # [B 2] + datacube["gsd"], # 1 + datacube["waves"], # [N] + ) # [B C H W] + B, C, H, W = cube.shape + + patches, _ = self.to_patch_embed( + cube, waves + ) # [B L D] - patchify & create embeddings per patch + + # Add time & latlon as encoding to patches + patches = self.add_encodings( + patches, + time, + latlon, + gsd, + ) # [B L D] - add position encoding to the embeddings + + # Add class tokens + cls_tokens = repeat(self.cls_token, "1 1 D -> B 1 D", B=B) # [B 1 D] + patches = torch.cat((cls_tokens, patches), dim=1) # [B (1 + L) D] + + # pass the patches through the transformer + patches = self.transformer(patches) # [B (1 + L) D] + + # get the cls token + embeddings = patches[:, 0, :] # [B D] + + return embeddings + + +class Embedder(nn.Module): + def __init__(self, img_size=256, ckpt_path=None, device="cpu"): + super().__init__() + self.clay_encoder = ( + EmbeddingEncoder( # Default parameters for the Clay base model + img_size=img_size, + patch_size=8, + dim=768, + depth=12, + heads=12, + dim_head=64, + mlp_ratio=4.0, + ).to(device) + ) + self.img_size = img_size + self.device = torch.device(device) + self.load_clay_weights(ckpt_path) + + def load_clay_weights(self, ckpt_path): + "Load the weights from the Clay model encoder." + ckpt = torch.load(ckpt_path, map_location=self.device) + state_dict = ckpt.get("state_dict") + state_dict = { + re.sub(r"^model\.encoder\.", "", name): param + for name, param in state_dict.items() + if name.startswith("model.encoder") + } + + with torch.no_grad(): + for name, param in self.clay_encoder.named_parameters(): + if name in state_dict and param.size() == state_dict[name].size(): + param.data.copy_(state_dict[name]) # Copy the weights + else: + print(f"No matching parameter for {name} with size {param.size()}") + + for param in self.clay_encoder.parameters(): + param.requires_grad = False + + self.clay_encoder.eval() + + def forward(self, datacube): + embeddings = self.clay_encoder(datacube) + + return embeddings + + def fake_datacube(self): + "Generate a fake datacube for model export." + dummy_datacube = { + "pixels": torch.randn(2, 3, self.img_size, self.img_size), + "time": torch.randn(2, 4), + "latlon": torch.randn(2, 4), + "waves": torch.randn(3), + "gsd": torch.randn(1), + } + dummy_datacube = {k: v.to(self.device) for k, v in dummy_datacube.items()} + return dummy_datacube + + def export_to_onnx(self, name): + "Save the model to ONNX format." + + datacube = self.fake_datacube() + export_options = torch.onnx.ExportOptions(dynamic_shapes=True) + + # Export the model to ONNX format + onnx_program = torch.onnx.dynamo_export( + self.eval(), datacube, export_options=export_options + ) + + # Save the exported model + onnx_program.save(f"checkpoints/compiled/{name}") + print(f"Model exported to ONNX format: checkpoints/compiled/{name}") + + return onnx_program + + def export_to_torchep(self, name): + "Save the model to pytorch ExportedProgram format." + + datacube = self.fake_datacube() + + # dynamic shapes for model export + batch_size = Dim("batch_size", min=2, max=1000) + channel_bands = Dim("channel_bands", min=1, max=10) + dynamic_shapes = { + "datacube": { + "pixels": {0: batch_size, 1: channel_bands}, + "time": {0: batch_size}, + "latlon": {0: batch_size}, + "waves": {0: channel_bands}, + "gsd": {0: None}, + } + } + + # Export the model to pytorch ExportedProgram format + ep = torch.export.export( + self.eval(), + (datacube,), + dynamic_shapes=dynamic_shapes, + strict=True, + ) + + # Save the exported model + torch.export.save(ep, f"checkpoints/compiled/{name}") + print( + f"Model exported to pytorch ExportedProgram format: checkpoints/compiled/{name}" # noqa: E501 + ) + + return ep + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Export the Clay model.") + parser.add_argument( + "--img_size", + type=int, + default=256, + help="Image size for the model", + ) + parser.add_argument( + "--ckpt_path", + type=str, + default="checkpoints/clay-v1-base.ckpt", + help="Path to the Clay model checkpoint", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + help="Device to use for the model", + ) + parser.add_argument( + "--name", + type=str, + default="clay-base.pt", + help="Name of the exported model", + ) + parser.add_argument( + "--onnx", + action="store_true", + help="Export the model to ONNX format", + ) + parser.add_argument( + "--ep", + action="store_true", + help="Export the model to pytorch ExportedProgram format", + ) + + args = parser.parse_args() + + Path("checkpoints/compiled").mkdir(parents=True, exist_ok=True) + embedder = Embedder( + img_size=args.img_size, + ckpt_path=args.ckpt_path, + device=args.device, + ) + + if args.onnx: + embedder.export_to_onnx(args.name) + elif args.ep: + embedder.export_to_torchep(args.name) + else: + print("Please specify the format to export the model.") + parser.print_help() From 1f2fcc9a0c74518907bb2a41c9cba86ca97e5f29 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 13:16:38 +0530 Subject: [PATCH 05/83] Remove files from src, fix utils to run everything on same device --- src/benchmark_encoder.py | 80 ------------------------------------- src/export.py | 73 ---------------------------------- src/test_encoder.py | 86 ---------------------------------------- src/utils.py | 16 ++++---- 4 files changed, 7 insertions(+), 248 deletions(-) delete mode 100644 src/benchmark_encoder.py delete mode 100644 src/export.py delete mode 100644 src/test_encoder.py diff --git a/src/benchmark_encoder.py b/src/benchmark_encoder.py deleted file mode 100644 index 08c2b0ac..00000000 --- a/src/benchmark_encoder.py +++ /dev/null @@ -1,80 +0,0 @@ -import argparse -import time -import warnings - -import torch - -warnings.filterwarnings("ignore") - -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - -def get_data(): - """ - Generate random data tensors for model input. - """ - cube = torch.randn(128, 3, 256, 256).to(DEVICE) - timestep = torch.randn(128, 4).to(DEVICE) - latlon = torch.randn(128, 4).to(DEVICE) - waves = torch.randn(3).to(DEVICE) - gsd = torch.randn(1).to(DEVICE) - return cube, timestep, latlon, waves, gsd - - -def load_exported_model(eager=True): - """ - Load the exported model from a file. - - Args: - eager (bool): Flag to decide whether to use eager mode or compiled mode. - """ - print("Loading exported model") - ep = torch.export.load("checkpoints/compiled/encoder.pt") - if eager: - model = ep.module() - else: - model = torch.compile(ep.module(), backend="inductor") - return model - - -def benchmark_model(model): - """ - Benchmark the model by running inference on randomly generated data. - - Args: - model: The model to benchmark. - """ - print("Benchmarking model") - start = time.time() - for i in range(20): - cube, timestep, latlon, waves, gsd = get_data() - with torch.inference_mode(): - out = model(cube, timestep, latlon, waves, gsd) - print( - f"Iteration {i}: Output shapes - {out[0].shape}, {out[1].shape}, {out[2].shape}, {out[3].shape}" # noqa E501 - ) - print("Time taken for inference: ", time.time() - start) - - -def run(eager=True): - """ - Run the exported model and benchmark it. - - Args: - eager (bool): Flag to decide whether to use eager mode or compiled mode. - """ - print("Running model") - model = load_exported_model(eager=eager) - benchmark_model(model) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="Run benchmark for the exported model." - ) - parser.add_argument( - "--eager", action="store_true", help="Use eager mode for running the model." - ) - args = parser.parse_args() - - run(args.eager) diff --git a/src/export.py b/src/export.py deleted file mode 100644 index 70a65f1e..00000000 --- a/src/export.py +++ /dev/null @@ -1,73 +0,0 @@ -import warnings -from pathlib import Path - -import torch -from torch.export import Dim - -from src.model import ClayMAEModule - -warnings.filterwarnings("ignore") - -CHECKPOINT_PATH = "checkpoints/clay-v1-base.ckpt" -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -CHIP_SIZE = 256 - - -def get_data(): - """ - Generate random data tensors for model input. - """ - cube = torch.randn(128, 3, CHIP_SIZE, CHIP_SIZE).to(DEVICE) - timestep = torch.randn(128, 4).to(DEVICE) - latlon = torch.randn(128, 4).to(DEVICE) - waves = torch.randn(3).to(DEVICE) - gsd = torch.randn(1).to(DEVICE) - return cube, timestep, latlon, waves, gsd - - -def load_model(): - """ - Load the model from a checkpoint and prepare it for evaluation. - """ - module = ClayMAEModule.load_from_checkpoint( - CHECKPOINT_PATH, shuffle=False, mask_ratio=0.0 - ) - encoder = module.model.encoder.eval() # Get the encoder in eval mode - encoder = encoder.to(DEVICE) # Move to the appropriate device - return encoder - - -def export_model(): - """ - Export the model with dynamic shapes for deployment. - """ - cube, timestep, latlon, waves, gsd = get_data() - encoder = load_model() - - # Define dynamic shapes for model export - batch_size = Dim("batch_size", min=32, max=1200) - channel_bands = Dim("channel_bands", min=1, max=10) - - dynamic_shapes = { - "cube": {0: batch_size, 1: channel_bands}, - "time": {0: batch_size}, - "latlon": {0: batch_size}, - "waves": {0: channel_bands}, - "gsd": {0: None}, - } - - # Export model - ep = torch.export.export( - mod=encoder, - args=(cube, timestep, latlon, waves, gsd), - dynamic_shapes=dynamic_shapes, - strict=True, - ) - - # Save the exported model - Path("checkpoints/compiled").mkdir(parents=True, exist_ok=True) - torch.export.save(ep, "checkpoints/compiled/encoder.pt") - - -if __name__ == "__main__": - export_model() diff --git a/src/test_encoder.py b/src/test_encoder.py deleted file mode 100644 index 14839197..00000000 --- a/src/test_encoder.py +++ /dev/null @@ -1,86 +0,0 @@ -import torch - -from src.datamodule import ClayDataModule - -# Load the pre-trained Clay encoder model -clay_encoder = torch.export.load("checkpoints/compiled/encoder.pt").module() - - -def load_batch(): - # Initialize the data module with appropriate parameters - dm = ClayDataModule( - data_dir="/home/ubuntu/data", - size=256, - metadata_path="configs/metadata.yaml", - batch_size=1, - num_workers=1, - ) - - # Setup the data module for the 'fit' stage - dm.setup(stage="fit") - metadata = dm.metadata - - # Get the training data loader and create an iterator - trn_dl = dm.train_dataloader() - iter_dl = iter(trn_dl) - - return iter_dl, metadata - - -def prepare_data(sensor, metadata, device): - """ - Load data from the sensor and transfer it to the specified device. - - Args: - - sensor (dict): Sensor data containing 'pixels', 'time', 'latlon', and 'platform'. - - metadata (dict): Metadata information for different platforms. - - device (torch.device): The device to which the data should be transferred. - - Returns: - - tuple: Transferred cube, timestep, latlon, waves, and gsd tensors. - """ - cube = sensor["pixels"] - timestep = sensor["time"] - latlon = sensor["latlon"] - platform = sensor["platform"][0] - - # Get wavelengths and ground sampling distance (gsd) from metadata - waves = torch.tensor(list(metadata[platform].bands.wavelength.values())) - gsd = torch.tensor([metadata[platform].gsd]) - - # Transfer data to the specified device - cube, timestep, latlon, waves, gsd = map( - lambda x: x.to(device), (cube, timestep, latlon, waves, gsd) - ) - return cube, timestep, latlon, waves, gsd - - -def main(): - dl, metadata = load_batch() - - # Fetch samples from the data loader - l8_c2l1 = next(dl) - l8_c2l2 = next(dl) - linz = next(dl) - naip = next(dl) - s1 = next(dl) - s2 = next(dl) - - # Perform inference with the Clay encoder model - with torch.no_grad(): - for sensor in (l8_c2l1, l8_c2l2, linz, naip, s1, s2): - # Load data and transfer to GPU - batch = prepare_data(sensor, metadata, torch.device("cuda")) - - # Get patch embeddings from the encoder model - patch_embeddings, *_ = clay_encoder(*batch) - - # Extract the class (CLS) embedding - cls_embedding = patch_embeddings[:, 0, :] - - # Print the platform and the shape of the CLS embedding - print(sensor["platform"][0], cls_embedding.shape) - - -if __name__ == "__main__": - main() diff --git a/src/utils.py b/src/utils.py index 1e35731f..b0f2bcce 100644 --- a/src/utils.py +++ b/src/utils.py @@ -11,7 +11,6 @@ def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype=torch.float32): assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb" omega = torch.arange(dim // 4) / (dim // 4 - 1) omega = 1.0 / (temperature**omega) - omega = omega.to(y.device) y = y.flatten()[:, None] * omega[None, :] x = x.flatten()[:, None] * omega[None, :] @@ -25,9 +24,9 @@ def posemb_sincos_2d_with_gsd( y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb" - omega = torch.arange(dim // 4, device=gsd.device) / (dim // 4 - 1) + gsd = gsd.to(x.device) + omega = torch.arange(dim // 4) / (dim // 4 - 1) omega = 1.0 / (temperature ** (2 * omega / dim)) * (gsd / 1.0) # Adjusted for g - omega = omega.to(y.device) y = y.flatten()[:, None] * omega[None, :] x = x.flatten()[:, None] * omega[None, :] @@ -35,17 +34,16 @@ def posemb_sincos_2d_with_gsd( return pe.type(dtype) -def posemb_sincos_1d(pos, dim, temperature: int = 10000, dtype=torch.float32): +def posemb_sincos_1d(waves, dim, temperature: int = 10000, dtype=torch.float32): assert ( dim % 2 == 0 ), "Feature dimension must be a multiple of 2 for sincos embedding" - pos = torch.arange(pos) if isinstance(pos, int) else pos + waves = torch.arange(waves) if isinstance(waves, int) else waves - omega = torch.arange(dim // 2) / (dim // 2 - 1) + omega = torch.arange(dim // 2, device=waves.device) / (dim // 2 - 1) omega = 1.0 / (temperature**omega) - omega = omega.to(pos.device) - scaled_pos = pos[:, None] * omega[None, :] - pe = torch.cat((scaled_pos.sin(), scaled_pos.cos()), dim=1) + scaled_waves = waves[:, None] * omega[None, :] + pe = torch.cat((scaled_waves.sin(), scaled_waves.cos()), dim=1) return pe.type(dtype) From 0a3ce9d90e4e9f2871b7d9db26b06d2583aef6bd Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 14:07:42 +0530 Subject: [PATCH 06/83] Bump torch==2.3.1 & torchvision==0.18.1, add onnx & onnxsxript as dependency --- environment.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/environment.yml b/environment.yml index 699df9df..2190062f 100644 --- a/environment.yml +++ b/environment.yml @@ -14,10 +14,12 @@ dependencies: - lancedb~=0.10.2 - lightning~=2.1.0 - matplotlib-base~=3.8.2 + - onnx~=1.16.1 + - onnxscript~=0.1.0.dev20240724 - planetary-computer~=1.0.0 - python-box~=7.1.0 - pytorch~=2.1.0 # [osx] - - pytorch~=2.1.0 *cuda12* # [linux] + - pytorch~=2.3.1 *cuda12* # [linux] - python~=3.11.0 - pyarrow~=16.1.0 - rioxarray~=0.15.0 @@ -29,7 +31,7 @@ dependencies: - timm~=0.9.16 - torchdata~=0.7.1 - torchgeo~=0.5.2 - - torchvision~=0.16.1 + - torchvision~=0.18.1 - transformers~=4.35.2 - typeshed-client~=2.4.0 - vit-pytorch~=1.6.4 From 37503feb487e377021ae62a0ee3386e5f976469e Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 16:40:03 +0530 Subject: [PATCH 07/83] Release few contraints on env --- environment.yml | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/environment.yml b/environment.yml index 2190062f..79a56ec8 100644 --- a/environment.yml +++ b/environment.yml @@ -11,14 +11,13 @@ dependencies: - jupyter-book~=1.0.0 - jupyterlab~=4.0.7 - jsonargparse~=4.27.0 - - lancedb~=0.10.2 - lightning~=2.1.0 - matplotlib-base~=3.8.2 - onnx~=1.16.1 - - onnxscript~=0.1.0.dev20240724 + - onnxscript - planetary-computer~=1.0.0 - python-box~=7.1.0 - - pytorch~=2.1.0 # [osx] + - pytorch~=2.3.1 # [osx] - pytorch~=2.3.1 *cuda12* # [linux] - python~=3.11.0 - pyarrow~=16.1.0 @@ -29,13 +28,12 @@ dependencies: - scikit-learn~=1.4.0 - stackstac~=0.5.0 - timm~=0.9.16 - - torchdata~=0.7.1 - - torchgeo~=0.5.2 + - torchgeo - torchvision~=0.18.1 - transformers~=4.35.2 - typeshed-client~=2.4.0 - vit-pytorch~=1.6.4 - - wandb~=0.15.12 + - wandb - zarr~=2.16.1 platforms: - linux-64 From db8a3f2a033835606e1360f6efdec9130599ae4b Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 12:52:26 +0000 Subject: [PATCH 08/83] Add notebook to show how to embed using compiled embedders --- environment.yml | 18 +- finetune/embedder/how-to-embed.ipynb | 637 +++++++++++++++++++++++++++ 2 files changed, 647 insertions(+), 8 deletions(-) create mode 100644 finetune/embedder/how-to-embed.ipynb diff --git a/environment.yml b/environment.yml index 79a56ec8..2eb2d1c3 100644 --- a/environment.yml +++ b/environment.yml @@ -7,33 +7,35 @@ dependencies: - einops~=0.7.0 - fiona~=1.9.5 - geopandas-base~=0.14.1 - - h5netcdf~=1.3.0 - - jupyter-book~=1.0.0 - - jupyterlab~=4.0.7 - jsonargparse~=4.27.0 - lightning~=2.1.0 - matplotlib-base~=3.8.2 - - onnx~=1.16.1 - - onnxscript - planetary-computer~=1.0.0 - python-box~=7.1.0 - pytorch~=2.3.1 # [osx] - pytorch~=2.3.1 *cuda12* # [linux] - python~=3.11.0 - pyarrow~=16.1.0 - - rioxarray~=0.15.0 - rasterio~=1.3.10 - s3fs~=2024.3.1 - scikit-image~=0.22.0 - scikit-learn~=1.4.0 - stackstac~=0.5.0 - timm~=0.9.16 - - torchgeo - torchvision~=0.18.1 - transformers~=4.35.2 - typeshed-client~=2.4.0 - vit-pytorch~=1.6.4 - - wandb - zarr~=2.16.1 + - pip: + - geoarrow-pyarrow==0.1.2 + - jupyter-book==1.0.2 + - jupyterlab==4.2.4 + - onnx==1.16.1 + - onnxscript + - onnxruntime + - torchgeo==0.5.2 + - stacchip==0.1.35 + - wandb==0.17.5 platforms: - linux-64 diff --git a/finetune/embedder/how-to-embed.ipynb b/finetune/embedder/how-to-embed.ipynb new file mode 100644 index 00000000..06f55cc7 --- /dev/null +++ b/finetune/embedder/how-to-embed.ipynb @@ -0,0 +1,637 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "d9960547-640d-425c-8180-fc5523a80e42", + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "import os\n", + "import requests\n", + "import warnings\n", + "\n", + "import geoarrow.pyarrow as ga\n", + "import numpy as np\n", + "import pystac_client\n", + "import pyarrow as pa\n", + "import pyarrow.parquet as pq\n", + "import torch\n", + "import yaml\n", + "from box import Box\n", + "from torchvision.transforms import v2\n", + "\n", + "from stacchip.indexer import Sentinel2Indexer\n", + "from stacchip.chipper import Chipper\n", + "\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "markdown", + "id": "598fec81-2cc1-4c5a-9e46-7c46a5591484", + "metadata": {}, + "source": [ + "### Find data for AOI\n", + "The first step is to find STAC items of imagery that we want to use to create embeddings. In this example we are going to use Earth Genome's composite dataset which comes with a great STAC catalog.\n", + "\n", + "We are also going to create embeddings along time so that we have multiple embeddings for the same location at different moments in time." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "3e1d46ee-40f6-49f5-99ad-83819339561e", + "metadata": {}, + "outputs": [], + "source": [ + "# Point over Monchique Portugal\n", + "lat, lon = 37.30939, -8.57207\n", + "\n", + "# Dates of a large forest fire\n", + "start = \"2018-07-01\"\n", + "end = \"2018-09-01\"" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "b7825318-23f3-449f-9104-eae6562a55ab", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found 12 items\n" + ] + } + ], + "source": [ + "# Optimize GDAL settings for cloud optimized reading\n", + "os.environ[\"GDAL_DISABLE_READDIR_ON_OPEN\"] = \"EMPTY_DIR\"\n", + "os.environ[\"AWS_REQUEST_PAYER\"] = \"requester\"\n", + "\n", + "STAC_API = \"https://earth-search.aws.element84.com/v1\"\n", + "COLLECTION = \"sentinel-2-l2a\"\n", + "\n", + "# Search the catalogue\n", + "catalog = pystac_client.Client.open(STAC_API)\n", + "search = catalog.search(\n", + " collections=[COLLECTION],\n", + " datetime=f\"{start}/{end}\",\n", + " bbox=(lon - 1e-5, lat - 1e-5, lon + 1e-5, lat + 1e-5),\n", + " max_items=100,\n", + " query={\"eo:cloud_cover\": {\"lt\": 80}},\n", + ")\n", + "\n", + "all_items = search.get_all_items()\n", + "\n", + "# Reduce to one per date (there might be some duplicates\n", + "# based on the location)\n", + "items = []\n", + "dates = []\n", + "for item in all_items:\n", + " if item.datetime.date() not in dates:\n", + " items.append(item)\n", + " dates.append(item.datetime.date())\n", + "\n", + "print(f\"Found {len(items)} items\")" + ] + }, + { + "cell_type": "markdown", + "id": "600f3cfb-ce4e-4409-ae15-20f3a7107a62", + "metadata": {}, + "source": [ + "To speed up processing in this example, we limit the number of chips to 3 per Sentinel-2 scene. Remove this limit in a real use case." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "183975c7-8afb-49ef-8e70-790265719aea", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n", + "Working on \n" + ] + } + ], + "source": [ + "chips = []\n", + "datetimes = []\n", + "bboxs = []\n", + "chip_ids = []\n", + "item_ids = []\n", + "\n", + "for item in items:\n", + " print(f\"Working on {item}\")\n", + "\n", + " # Index the chips in the item\n", + " indexer = Sentinel2Indexer(item)\n", + "\n", + " # Instanciate the chipper\n", + " chipper = Chipper(indexer, assets=[\"red\", \"green\", \"blue\", \"nir\", \"scl\"])\n", + "\n", + " # Get first chip for the \"image\" asset key\n", + " for idx, (x, y, chip) in enumerate(chipper):\n", + " if idx > 2:\n", + " break\n", + " del chip[\"scl\"]\n", + " chips.append(chip)\n", + " datetimes.append(item.datetime)\n", + " bboxs.append(indexer.get_chip_bbox(x, y))\n", + " chip_ids.append((x, y))\n", + " item_ids.append(item.id)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "71902ab7-3320-43cd-85c3-362c2500f241", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "(36, 4, 256, 256)" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips])\n", + "pixels.shape" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "6f7ce367-4e12-4648-bb79-119b4f50ead8", + "metadata": {}, + "outputs": [], + "source": [ + "# Extract mean, std, and wavelengths from metadata\n", + "platform = \"sentinel-2-l2a\"\n", + "# Retrieve the file content from the URL\n", + "\n", + "url = (\n", + " \"https://raw.githubusercontent.com/Clay-foundation/model/main/configs/metadata.yaml\"\n", + ")\n", + "response = requests.get(url, allow_redirects=True)\n", + "\n", + "# Convert bytes to string\n", + "content = response.content.decode(\"utf-8\")\n", + "\n", + "# Load the yaml\n", + "content = yaml.safe_load(content)\n", + "\n", + "metadata = Box(content)\n", + "mean = []\n", + "std = []\n", + "waves = []\n", + "# Use the band names to get the correct values in the correct order.\n", + "for band in chips[0].keys():\n", + " mean.append(metadata[platform].bands.mean[band])\n", + " std.append(metadata[platform].bands.std[band])\n", + " waves.append(metadata[platform].bands.wavelength[band])\n", + "\n", + "# Prepare the normalization transform function using the mean and std values.\n", + "transform = v2.Compose(\n", + " [\n", + " v2.Normalize(mean=mean, std=std),\n", + " ]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "a8ec8c2d-ecb9-42a2-9e8c-3f95c67ef07b", + "metadata": {}, + "outputs": [], + "source": [ + "def normalize_timestamp(date):\n", + " week = date.isocalendar().week * 2 * np.pi / 52\n", + " hour = date.hour * 2 * np.pi / 24\n", + "\n", + " return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour))\n", + "\n", + "\n", + "times = [normalize_timestamp(dat) for dat in datetimes]\n", + "week_norm = [dat[0] for dat in times]\n", + "hour_norm = [dat[1] for dat in times]\n", + "\n", + "\n", + "# Prep lat/lon embedding using the\n", + "def normalize_latlon(lat, lon):\n", + " lat = lat * np.pi / 180\n", + " lon = lon * np.pi / 180\n", + "\n", + " return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon))\n", + "\n", + "\n", + "latlons = [normalize_latlon(lat, lon)] * len(times)\n", + "lat_norm = [dat[0] for dat in latlons]\n", + "lon_norm = [dat[1] for dat in latlons]\n", + "\n", + "# Prep gsd\n", + "gsd = [10]\n", + "\n", + "# Normalize pixels\n", + "pixels = transform(pixels)" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "2640eb17-a85c-4972-8d5d-e45e9ed8eba5", + "metadata": {}, + "outputs": [], + "source": [ + "datacube = {\n", + " \"pixels\": torch.tensor(pixels, dtype=torch.float32),\n", + " \"time\": torch.tensor(np.hstack((week_norm, hour_norm)), dtype=torch.float32),\n", + " \"latlon\": torch.tensor(np.hstack((lat_norm, lon_norm)), dtype=torch.float32),\n", + " \"waves\": torch.tensor(waves, dtype=torch.float32),\n", + " \"gsd\": torch.tensor(gsd, dtype=torch.float32),\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "7f6711a9-e7ed-44d5-add7-2c3a498cd422", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "pixels torch.Size([36, 4, 256, 256])\n", + "time torch.Size([36, 4])\n", + "latlon torch.Size([36, 4])\n", + "waves torch.Size([4])\n", + "gsd torch.Size([1])\n" + ] + } + ], + "source": [ + "for k,v in datacube.items():\n", + " print(k, v.shape)" + ] + }, + { + "cell_type": "markdown", + "id": "83243912-a2a8-4fa5-a39c-a9c3b07c7569", + "metadata": {}, + "source": [ + "### Clay Embedder\n", + "\n", + "#### Load the embedder that is stored in ExportedProgram format using **cpu**." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "4eb468af-d468-46aa-a8fb-23ff95c56288", + "metadata": {}, + "outputs": [], + "source": [ + "!wget -q https://huggingface.co/made-with-clay/Clay/resolve/main/compiled/v1.0/clay-v1-encoder-cpu.pt2" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "9eb797f7-5238-49e0-9950-e85f10132454", + "metadata": {}, + "outputs": [], + "source": [ + "ep_embedder_cpu = torch.export.load(\"clay-v1-encoder-cpu.pt2\").module()" + ] + }, + { + "cell_type": "code", + "execution_count": 19, + "id": "eefe4811-7290-47c3-a10e-45257e6d42e0", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 2min 36s, sys: 26.9 s, total: 3min 3s\n", + "Wall time: 51.3 s\n" + ] + }, + { + "data": { + "text/plain": [ + "(torch.Size([36, 4, 256, 256]), torch.Size([36, 768]))" + ] + }, + "execution_count": 19, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "with torch.no_grad():\n", + " embeddings = ep_embedder_cpu(datacube)\n", + "datacube[\"pixels\"].shape, embeddings.shape" + ] + }, + { + "cell_type": "markdown", + "id": "8e927b01-c855-4172-a4d9-2c10ba794ed4", + "metadata": {}, + "source": [ + "For each chip, we have an embedding of size `768`" + ] + }, + { + "cell_type": "markdown", + "id": "fa0810b4-34ad-490e-bbcd-c0c3288f017c", + "metadata": {}, + "source": [ + "#### Load the embedder that is stored in ExportedProgram format using **gpu**." + ] + }, + { + "cell_type": "code", + "execution_count": 21, + "id": "9c1bbfd4-7dc6-4ad0-8a0b-b3745a9f35ca", + "metadata": {}, + "outputs": [], + "source": [ + "!wget -q https://huggingface.co/made-with-clay/Clay/resolve/main/compiled/v1.0/clay-v1-encoder.pt2" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "e285a543-20ab-44ba-b676-2303284dc477", + "metadata": {}, + "outputs": [], + "source": [ + "datacube = {k:v.to(\"cuda\") for k,v in datacube.items()}\n", + "ep_embedder = torch.export.load(\"clay-v1-encoder.pt2\").module()" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "edefee90-e6b8-4701-bb5d-2bf7febc806c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 313 ms, sys: 41.5 ms, total: 354 ms\n", + "Wall time: 239 ms\n" + ] + }, + { + "data": { + "text/plain": [ + "(torch.Size([36, 4, 256, 256]), torch.Size([36, 768]))" + ] + }, + "execution_count": 23, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "with torch.no_grad():\n", + " embeddings = ep_embedder(datacube)\n", + "datacube[\"pixels\"].shape, embeddings.shape" + ] + }, + { + "cell_type": "markdown", + "id": "196f2121-46b5-4b02-94d3-75e648c329c3", + "metadata": {}, + "source": [ + "For each chip, we have an embedding of size `768`" + ] + }, + { + "cell_type": "markdown", + "id": "5b1cb0f9-a434-419b-a88b-4d4edd84fea6", + "metadata": {}, + "source": [ + "#### Load the embedder that is stored in ONNX format using **cpu**." + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "aa10d696-740a-458e-ae10-eec9a43fb362", + "metadata": {}, + "outputs": [], + "source": [ + "import onnx\n", + "import onnxruntime as ort" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "992524e5-2c2a-4e48-ae95-bd2aa87b72a9", + "metadata": {}, + "outputs": [], + "source": [ + "!wget -q https://huggingface.co/made-with-clay/Clay/resolve/main/compiled/v1.0/clay-v1-encoder-cpu.onnx" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "dc3fa967-73d5-431c-88a2-84b088aff06f", + "metadata": {}, + "outputs": [], + "source": [ + "datacube = {k:v.to(\"cpu\") for k,v in datacube.items()}\n", + "onnx_embedder = ort.InferenceSession(\"clay-v1-encoder-cpu.onnx\", \n", + " providers=[\"CPUExecutionProvider\"])" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "24591d17-d1c8-452b-9b20-676a9b6f8643", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 3min 48s, sys: 1.82 s, total: 3min 50s\n", + "Wall time: 30.6 s\n" + ] + }, + { + "data": { + "text/plain": [ + "(36, 768)" + ] + }, + "execution_count": 27, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "%%time\n", + "embeddings = onnx_embedder.run([], {\n", + " \"cube\": datacube[\"pixels\"].numpy(),\n", + " \"time\": datacube[\"time\"].numpy(),\n", + " \"latlon\": datacube[\"latlon\"].numpy(),\n", + " \"waves\": datacube[\"waves\"].numpy(),\n", + " \"gsd\": datacube[\"gsd\"].numpy()\n", + "})[0]\n", + "embeddings.shape" + ] + }, + { + "cell_type": "markdown", + "id": "9c07216e-a109-4cd8-8c74-9a3fc9a37757", + "metadata": {}, + "source": [ + "For each chip, we have an embedding of size `768`" + ] + }, + { + "cell_type": "markdown", + "id": "2e8d5900-9a4b-4e2d-b992-4fb0a1e8c835", + "metadata": {}, + "source": [ + "### Store the results\n", + "\n", + "We create a table containing the embeddings, bounding box, the STAC item ID, the datetime of the image capture, and the chip x and y ids. Then we save that data to disk." + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "677f04d3-db38-4d44-9b55-c103d54adcd5", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "pyarrow.Table\n", + "datetimes: timestamp[us, tz=UTC]\n", + "chip_ids: list\n", + " child 0, item: int64\n", + "item_ids: string\n", + "emeddings: list\n", + " child 0, item: float\n", + "geometry: extension>\n", + "----\n", + "datetimes: [[2018-08-28 11:30:56.771000Z,2018-08-28 11:30:56.771000Z,2018-08-28 11:30:56.771000Z,2018-08-23 11:30:50.574000Z,2018-08-23 11:30:50.574000Z,...,2018-07-09 11:24:55.535000Z,2018-07-09 11:24:55.535000Z,2018-07-04 11:30:35.271000Z,2018-07-04 11:30:35.271000Z,2018-07-04 11:30:35.271000Z]]\n", + "chip_ids: [[[0,0],[1,0],...,[1,0],[2,0]]]\n", + "item_ids: [[\"S2A_29SNB_20180828_1_L2A\",\"S2A_29SNB_20180828_1_L2A\",\"S2A_29SNB_20180828_1_L2A\",\"S2B_29SNB_20180823_1_L2A\",\"S2B_29SNB_20180823_1_L2A\",...,\"S2A_29SNB_20180709_0_L2A\",\"S2A_29SNB_20180709_0_L2A\",\"S2B_29SNB_20180704_0_L2A\",\"S2B_29SNB_20180704_0_L2A\",\"S2B_29SNB_20180704_0_L2A\"]]\n", + "emeddings: [[[-0.14773342,0.08466571,0.13797832,0.11150883,0.06517959,...,0.036681578,-0.092160255,0.025934512,-0.12496276,-0.034070153],[-0.14430065,0.085857555,0.13839196,0.10963549,0.0652737,...,0.03711322,-0.09153629,0.02631686,-0.12422915,-0.03333628],...,[-0.09626354,0.062443394,0.24817112,0.012715777,0.043093704,...,0.011770063,-0.037860263,0.027813748,-0.11962952,-0.02246455],[-0.10004063,0.06320572,0.24851695,0.012129029,0.043350283,...,0.011444314,-0.03733269,0.027787287,-0.12139094,-0.021088997]]]\n", + "geometry: [[[ -- is_valid: all not null\n", + " -- child 0 type: double\n", + "[-8.825403979293151,-8.825730459265694,-9.000227209792856,-9.000227635454767,-8.825403979293151]\n", + " -- child 1 type: double\n", + "[37.947460030545635,37.809019655564406,37.809148556380286,37.947589571562965,37.947460030545635]],[ -- is_valid: all not null\n", + " -- child 0 type: double\n", + "[-8.650582567535476,-8.651235936821893,-8.825730459265694,-8.825403979293151,-8.650582567535476]\n", + " -- child 1 type: double\n", + "[37.94707073614538,37.80863228507305,37.809019655564406,37.947460030545635,37.94707073614538]],...,[ -- is_valid: all not null\n", + " -- child 0 type: double\n", + "[-8.650582567535476,-8.651235936821893,-8.825730459265694,-8.825403979293151,-8.650582567535476]\n", + " -- child 1 type: double\n", + "[37.94707073614538,37.80863228507305,37.809019655564406,37.947460030545635,37.94707073614538]],[ -- is_valid: all not null\n", + " -- child 0 type: double\n", + "[-8.475765647330832,-8.476745873271028,-8.651235936821893,-8.650582567535476,-8.475765647330832]\n", + " -- child 1 type: double\n", + "[37.94642170369997,37.80798646012822,37.80863228507305,37.94707073614538,37.94642170369997]]]]" + ] + }, + "execution_count": 28, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Write data to pyarrow table\n", + "index = {\n", + " \"datetimes\": datetimes,\n", + " \"chip_ids\": chip_ids,\n", + " \"item_ids\": item_ids,\n", + " \"emeddings\": [np.ascontiguousarray(dat) for dat in embeddings],\n", + " \"geometry\": ga.as_geoarrow([dat.wkt for dat in bboxs]),\n", + "}\n", + "table = pa.table(index)\n", + "table" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "id": "d62a9e8a-b4f9-491c-a437-6a164a9e74fe", + "metadata": {}, + "outputs": [], + "source": [ + "pq.write_table(table, \"embeddings.parquet\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d30fb8c7-d04d-453f-93f6-dc3599f1df15", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From c0552bda2c326c84fcda2a259f5621e1fa843747 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jul 2024 12:53:09 +0000 Subject: [PATCH 09/83] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- finetune/embedder/how-to-embed.ipynb | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/finetune/embedder/how-to-embed.ipynb b/finetune/embedder/how-to-embed.ipynb index 06f55cc7..3799c12a 100644 --- a/finetune/embedder/how-to-embed.ipynb +++ b/finetune/embedder/how-to-embed.ipynb @@ -296,7 +296,7 @@ } ], "source": [ - "for k,v in datacube.items():\n", + "for k, v in datacube.items():\n", " print(k, v.shape)" ] }, @@ -395,7 +395,7 @@ "metadata": {}, "outputs": [], "source": [ - "datacube = {k:v.to(\"cuda\") for k,v in datacube.items()}\n", + "datacube = {k: v.to(\"cuda\") for k, v in datacube.items()}\n", "ep_embedder = torch.export.load(\"clay-v1-encoder.pt2\").module()" ] }, @@ -475,9 +475,10 @@ "metadata": {}, "outputs": [], "source": [ - "datacube = {k:v.to(\"cpu\") for k,v in datacube.items()}\n", - "onnx_embedder = ort.InferenceSession(\"clay-v1-encoder-cpu.onnx\", \n", - " providers=[\"CPUExecutionProvider\"])" + "datacube = {k: v.to(\"cpu\") for k, v in datacube.items()}\n", + "onnx_embedder = ort.InferenceSession(\n", + " \"clay-v1-encoder-cpu.onnx\", providers=[\"CPUExecutionProvider\"]\n", + ")" ] }, { @@ -507,13 +508,16 @@ ], "source": [ "%%time\n", - "embeddings = onnx_embedder.run([], {\n", - " \"cube\": datacube[\"pixels\"].numpy(),\n", - " \"time\": datacube[\"time\"].numpy(),\n", - " \"latlon\": datacube[\"latlon\"].numpy(),\n", - " \"waves\": datacube[\"waves\"].numpy(),\n", - " \"gsd\": datacube[\"gsd\"].numpy()\n", - "})[0]\n", + "embeddings = onnx_embedder.run(\n", + " [],\n", + " {\n", + " \"cube\": datacube[\"pixels\"].numpy(),\n", + " \"time\": datacube[\"time\"].numpy(),\n", + " \"latlon\": datacube[\"latlon\"].numpy(),\n", + " \"waves\": datacube[\"waves\"].numpy(),\n", + " \"gsd\": datacube[\"gsd\"].numpy(),\n", + " },\n", + ")[0]\n", "embeddings.shape" ] }, From 1e9750661ff534b504600ce70c6c3677c97e0cc4 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 12:57:50 +0000 Subject: [PATCH 10/83] Clear outputs from the notebook --- finetune/embedder/how-to-embed.ipynb | 205 ++++----------------------- 1 file changed, 29 insertions(+), 176 deletions(-) diff --git a/finetune/embedder/how-to-embed.ipynb b/finetune/embedder/how-to-embed.ipynb index 06f55cc7..77e4ed2c 100644 --- a/finetune/embedder/how-to-embed.ipynb +++ b/finetune/embedder/how-to-embed.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "d9960547-640d-425c-8180-fc5523a80e42", "metadata": {}, "outputs": [], @@ -41,7 +41,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "3e1d46ee-40f6-49f5-99ad-83819339561e", "metadata": {}, "outputs": [], @@ -56,18 +56,10 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "b7825318-23f3-449f-9104-eae6562a55ab", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Found 12 items\n" - ] - } - ], + "outputs": [], "source": [ "# Optimize GDAL settings for cloud optimized reading\n", "os.environ[\"GDAL_DISABLE_READDIR_ON_OPEN\"] = \"EMPTY_DIR\"\n", @@ -110,29 +102,10 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "183975c7-8afb-49ef-8e70-790265719aea", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n", - "Working on \n" - ] - } - ], + "outputs": [], "source": [ "chips = []\n", "datetimes = []\n", @@ -163,21 +136,10 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "71902ab7-3320-43cd-85c3-362c2500f241", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(36, 4, 256, 256)" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips])\n", "pixels.shape" @@ -185,7 +147,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "6f7ce367-4e12-4648-bb79-119b4f50ead8", "metadata": {}, "outputs": [], @@ -225,7 +187,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "a8ec8c2d-ecb9-42a2-9e8c-3f95c67ef07b", "metadata": {}, "outputs": [], @@ -263,7 +225,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "id": "2640eb17-a85c-4972-8d5d-e45e9ed8eba5", "metadata": {}, "outputs": [], @@ -279,22 +241,10 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "id": "7f6711a9-e7ed-44d5-add7-2c3a498cd422", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "pixels torch.Size([36, 4, 256, 256])\n", - "time torch.Size([36, 4])\n", - "latlon torch.Size([36, 4])\n", - "waves torch.Size([4])\n", - "gsd torch.Size([1])\n" - ] - } - ], + "outputs": [], "source": [ "for k,v in datacube.items():\n", " print(k, v.shape)" @@ -312,7 +262,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "id": "4eb468af-d468-46aa-a8fb-23ff95c56288", "metadata": {}, "outputs": [], @@ -322,7 +272,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "id": "9eb797f7-5238-49e0-9950-e85f10132454", "metadata": {}, "outputs": [], @@ -332,29 +282,10 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "id": "eefe4811-7290-47c3-a10e-45257e6d42e0", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 2min 36s, sys: 26.9 s, total: 3min 3s\n", - "Wall time: 51.3 s\n" - ] - }, - { - "data": { - "text/plain": [ - "(torch.Size([36, 4, 256, 256]), torch.Size([36, 768]))" - ] - }, - "execution_count": 19, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "with torch.no_grad():\n", @@ -380,7 +311,7 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "id": "9c1bbfd4-7dc6-4ad0-8a0b-b3745a9f35ca", "metadata": {}, "outputs": [], @@ -390,7 +321,7 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "id": "e285a543-20ab-44ba-b676-2303284dc477", "metadata": {}, "outputs": [], @@ -401,29 +332,10 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "id": "edefee90-e6b8-4701-bb5d-2bf7febc806c", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 313 ms, sys: 41.5 ms, total: 354 ms\n", - "Wall time: 239 ms\n" - ] - }, - { - "data": { - "text/plain": [ - "(torch.Size([36, 4, 256, 256]), torch.Size([36, 768]))" - ] - }, - "execution_count": 23, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "with torch.no_grad():\n", @@ -449,7 +361,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": null, "id": "aa10d696-740a-458e-ae10-eec9a43fb362", "metadata": {}, "outputs": [], @@ -460,7 +372,7 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": null, "id": "992524e5-2c2a-4e48-ae95-bd2aa87b72a9", "metadata": {}, "outputs": [], @@ -470,7 +382,7 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "id": "dc3fa967-73d5-431c-88a2-84b088aff06f", "metadata": {}, "outputs": [], @@ -482,29 +394,10 @@ }, { "cell_type": "code", - "execution_count": 27, + "execution_count": null, "id": "24591d17-d1c8-452b-9b20-676a9b6f8643", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 3min 48s, sys: 1.82 s, total: 3min 50s\n", - "Wall time: 30.6 s\n" - ] - }, - { - "data": { - "text/plain": [ - "(36, 768)" - ] - }, - "execution_count": 27, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "embeddings = onnx_embedder.run([], {\n", @@ -537,50 +430,10 @@ }, { "cell_type": "code", - "execution_count": 28, + "execution_count": null, "id": "677f04d3-db38-4d44-9b55-c103d54adcd5", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "pyarrow.Table\n", - "datetimes: timestamp[us, tz=UTC]\n", - "chip_ids: list\n", - " child 0, item: int64\n", - "item_ids: string\n", - "emeddings: list\n", - " child 0, item: float\n", - "geometry: extension>\n", - "----\n", - "datetimes: [[2018-08-28 11:30:56.771000Z,2018-08-28 11:30:56.771000Z,2018-08-28 11:30:56.771000Z,2018-08-23 11:30:50.574000Z,2018-08-23 11:30:50.574000Z,...,2018-07-09 11:24:55.535000Z,2018-07-09 11:24:55.535000Z,2018-07-04 11:30:35.271000Z,2018-07-04 11:30:35.271000Z,2018-07-04 11:30:35.271000Z]]\n", - "chip_ids: [[[0,0],[1,0],...,[1,0],[2,0]]]\n", - "item_ids: [[\"S2A_29SNB_20180828_1_L2A\",\"S2A_29SNB_20180828_1_L2A\",\"S2A_29SNB_20180828_1_L2A\",\"S2B_29SNB_20180823_1_L2A\",\"S2B_29SNB_20180823_1_L2A\",...,\"S2A_29SNB_20180709_0_L2A\",\"S2A_29SNB_20180709_0_L2A\",\"S2B_29SNB_20180704_0_L2A\",\"S2B_29SNB_20180704_0_L2A\",\"S2B_29SNB_20180704_0_L2A\"]]\n", - "emeddings: [[[-0.14773342,0.08466571,0.13797832,0.11150883,0.06517959,...,0.036681578,-0.092160255,0.025934512,-0.12496276,-0.034070153],[-0.14430065,0.085857555,0.13839196,0.10963549,0.0652737,...,0.03711322,-0.09153629,0.02631686,-0.12422915,-0.03333628],...,[-0.09626354,0.062443394,0.24817112,0.012715777,0.043093704,...,0.011770063,-0.037860263,0.027813748,-0.11962952,-0.02246455],[-0.10004063,0.06320572,0.24851695,0.012129029,0.043350283,...,0.011444314,-0.03733269,0.027787287,-0.12139094,-0.021088997]]]\n", - "geometry: [[[ -- is_valid: all not null\n", - " -- child 0 type: double\n", - "[-8.825403979293151,-8.825730459265694,-9.000227209792856,-9.000227635454767,-8.825403979293151]\n", - " -- child 1 type: double\n", - "[37.947460030545635,37.809019655564406,37.809148556380286,37.947589571562965,37.947460030545635]],[ -- is_valid: all not null\n", - " -- child 0 type: double\n", - "[-8.650582567535476,-8.651235936821893,-8.825730459265694,-8.825403979293151,-8.650582567535476]\n", - " -- child 1 type: double\n", - "[37.94707073614538,37.80863228507305,37.809019655564406,37.947460030545635,37.94707073614538]],...,[ -- is_valid: all not null\n", - " -- child 0 type: double\n", - "[-8.650582567535476,-8.651235936821893,-8.825730459265694,-8.825403979293151,-8.650582567535476]\n", - " -- child 1 type: double\n", - "[37.94707073614538,37.80863228507305,37.809019655564406,37.947460030545635,37.94707073614538]],[ -- is_valid: all not null\n", - " -- child 0 type: double\n", - "[-8.475765647330832,-8.476745873271028,-8.651235936821893,-8.650582567535476,-8.475765647330832]\n", - " -- child 1 type: double\n", - "[37.94642170369997,37.80798646012822,37.80863228507305,37.94707073614538,37.94642170369997]]]]" - ] - }, - "execution_count": 28, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Write data to pyarrow table\n", "index = {\n", @@ -596,7 +449,7 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": null, "id": "d62a9e8a-b4f9-491c-a437-6a164a9e74fe", "metadata": {}, "outputs": [], From 1803954f73983ede18483c009361ec32ab7f89d7 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 25 Jul 2024 15:26:17 +0000 Subject: [PATCH 11/83] Add torchdata as a pip dependency --- environment.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/environment.yml b/environment.yml index 2eb2d1c3..ac0ddeaf 100644 --- a/environment.yml +++ b/environment.yml @@ -34,6 +34,7 @@ dependencies: - onnx==1.16.1 - onnxscript - onnxruntime + - torchdata==0.7.1 - torchgeo==0.5.2 - stacchip==0.1.35 - wandb==0.17.5 From fdf4e80242ae87d3495ce11e1eb5beb8bbef61ab Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 13:08:27 +0530 Subject: [PATCH 12/83] Randomly pass time & latlon as zeros 20% of time --- src/datamodule.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/src/datamodule.py b/src/datamodule.py index 7ba22e98..247df5b5 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -3,6 +3,7 @@ rasterio. """ +import random from collections import defaultdict from pathlib import Path from typing import List, Literal @@ -57,16 +58,23 @@ def __getitem__(self, idx): platform = chip_path.parent.name pixels = self.transforms[platform](pixels) + time_tensor = torch.tensor( + np.hstack((chip["week_norm"], chip["hour_norm"]), dtype=np.float32) + ) + latlon_tensor = torch.tensor( + np.hstack((chip["lat_norm"], chip["lon_norm"]), dtype=np.float32) + ) + + # Randomly set time & latlon to zero for 20% of the chips + if random.random() < 0.2: # noqa: PLR2004 + time_tensor.zero_() + latlon_tensor.zero_() + # Prepare additional information additional_info = { "platform": platform, - "time": torch.tensor( - np.hstack((chip["week_norm"], chip["hour_norm"])), - dtype=torch.float32, - ), - "latlon": torch.tensor( - np.hstack((chip["lat_norm"], chip["lon_norm"])), dtype=torch.float32 - ), + "time": time_tensor, + "latlon": latlon_tensor, } return {"pixels": pixels, **additional_info} From 519a171968c37aea7264897f0478df32028190eb Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 13:10:20 +0530 Subject: [PATCH 13/83] Add modis band info to metadata.yaml --- configs/metadata.yaml | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/configs/metadata.yaml b/configs/metadata.yaml index 7467b120..4309c5e2 100644 --- a/configs/metadata.yaml +++ b/configs/metadata.yaml @@ -184,3 +184,38 @@ sentinel-1-rtc: wavelength: vv: 3.5 vh: 4.0 +modis: + band_order: + - sur_refl_b01 + - sur_refl_b02 + - sur_refl_b03 + - sur_refl_b04 + - sur_refl_b05 + - sur_refl_b06 + - sur_refl_b07 + gsd: 500 + bands: + mean: + sur_refl_b01: 1072. + sur_refl_b02: 1624. + sur_refl_b03: 931. + sur_refl_b04: 1023. + sur_refl_b05: 1599. + sur_refl_b06: 1404. + sur_refl_b07: 1051. + std: + sur_refl_b01: 1643. + sur_refl_b02: 1878. + sur_refl_b03: 1449. + sur_refl_b04: 1538. + sur_refl_b05: 1763. + sur_refl_b06: 1618. + sur_refl_b07: 1396. + wavelength: + sur_refl_b01: 645. + sur_refl_b02: 858. + sur_refl_b03: 469. + sur_refl_b04: 555. + sur_refl_b05: 1240. + sur_refl_b06: 1640. + sur_refl_b07: 2130. From a9efd2b24c1b3b34814ac584c64a6f33ec1af499 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 13:12:47 +0530 Subject: [PATCH 14/83] Add prefetch factor as an arg to DataModule --- src/datamodule.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/datamodule.py b/src/datamodule.py index 247df5b5..7cfe579f 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -143,12 +143,14 @@ def __init__( # noqa: PLR0913 "landsat-c2l1", "landsat-c2l2-sr", "linz", + "modis", "naip", "sentinel-1-rtc", "sentinel-2-l2a", ], batch_size: int = 10, num_workers: int = 8, + prefetch_factor: int = 2, ): super().__init__() self.data_dir = data_dir @@ -157,6 +159,7 @@ def __init__( # noqa: PLR0913 self.metadata = Box(yaml.safe_load(open(metadata_path))) self.batch_size = batch_size self.num_workers = num_workers + self.prefetch_factor = prefetch_factor self.split_ratio = 0.8 def setup(self, stage: Literal["fit", "predict"] | None = None) -> None: @@ -215,7 +218,7 @@ def train_dataloader(self): batch_sampler=self.trn_sampler, collate_fn=batch_collate, pin_memory=True, - prefetch_factor=4, + prefetch_factor=self.prefetch_factor, ) def val_dataloader(self): @@ -225,7 +228,7 @@ def val_dataloader(self): batch_sampler=self.val_sampler, collate_fn=batch_collate, pin_memory=True, - prefetch_factor=4, + prefetch_factor=self.prefetch_factor, ) def predict_dataloader(self): From 6434473373a5b386b7a2950daefe06dfff985c59 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 16:40:41 +0530 Subject: [PATCH 15/83] Change wavelength to millimeter scale for modis --- configs/metadata.yaml | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/configs/metadata.yaml b/configs/metadata.yaml index 4309c5e2..7e3d71af 100644 --- a/configs/metadata.yaml +++ b/configs/metadata.yaml @@ -193,6 +193,10 @@ modis: - sur_refl_b05 - sur_refl_b06 - sur_refl_b07 + rgb_indices: + - 0 + - 3 + - 2 gsd: 500 bands: mean: @@ -212,10 +216,10 @@ modis: sur_refl_b06: 1618. sur_refl_b07: 1396. wavelength: - sur_refl_b01: 645. - sur_refl_b02: 858. - sur_refl_b03: 469. - sur_refl_b04: 555. - sur_refl_b05: 1240. - sur_refl_b06: 1640. - sur_refl_b07: 2130. + sur_refl_b01: .645 + sur_refl_b02: .858 + sur_refl_b03: .469 + sur_refl_b04: .555 + sur_refl_b05: 1.240 + sur_refl_b06: 1.640 + sur_refl_b07: 2.130 From 7021ab4d888d71e90fabdd71c47037c08cd90260 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 20:36:07 +0530 Subject: [PATCH 16/83] change batch_first to True --- src/factory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/factory.py b/src/factory.py index 9f10fab9..84a0d5b5 100644 --- a/src/factory.py +++ b/src/factory.py @@ -44,7 +44,7 @@ def __init__( # noqa: PLR0913 activation="gelu", dropout=0, norm_first=False, - batch_first=False, + batch_first=True, ) self.encoder = nn.TransformerEncoder(layer, num_layers) From 5770cd669603b098cffd3f52657fc22e27f3f7f0 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 20:38:47 +0530 Subject: [PATCH 17/83] Add transformer code from vit_pytorch as module --- src/backbone.py | 83 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100644 src/backbone.py diff --git a/src/backbone.py b/src/backbone.py new file mode 100644 index 00000000..16fd1fd5 --- /dev/null +++ b/src/backbone.py @@ -0,0 +1,83 @@ +"""Code for Transformer from Phil Wangs library +Repository: https://github.com/lucidrains/vit-pytorch +""" + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import nn + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim): + super().__init__() + self.net = nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, dim), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(nn.Module): + def __init__(self, dim, heads=8, dim_head=64, fused_attn=True): + super().__init__() + inner_dim = dim_head * heads + self.heads = heads + self.scale = dim_head**-0.5 + self.norm = nn.LayerNorm(dim) + self.fused_attn = fused_attn + + self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x): + x = self.norm(x) + + qkv = self.to_qkv(x).chunk(3, dim=-1) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv) + + if self.fused_attn: + x = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0) + else: + attn = torch.matmul(q, k.transpose(-1, -2)) * self.scale + attn = attn.softmax(dim=-1) + x = torch.matmul(attn, v) + + x = rearrange(x, "b h n d -> b n (h d)") + return self.to_out(x) + + +class Transformer(nn.Module): + def __init__( # noqa: PLR0913 + self, + dim, + depth, + heads, + dim_head, + mlp_dim, + fused_attn, + ): + super().__init__() + self.norm = nn.LayerNorm(dim) + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + Attention( + dim, heads=heads, dim_head=dim_head, fused_attn=fused_attn + ), + FeedForward(dim, mlp_dim), + ] + ) + ) + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + return self.norm(x) From 6bbc9026f6274d48a0933f2dfc994756434a33d3 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 20:40:48 +0530 Subject: [PATCH 18/83] Add MRL --- src/mrl.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 src/mrl.py diff --git a/src/mrl.py b/src/mrl.py new file mode 100644 index 00000000..a7202ad4 --- /dev/null +++ b/src/mrl.py @@ -0,0 +1,36 @@ +from torch import nn + + +class MRL(nn.Module): + """ + Matryoshka Representation Learning from the paper: https://arxiv.org/abs/2205.13147 + """ + + def __init__(self, features, dolls: list = [16, 32, 64, 128, 256, 768]) -> None: + super().__init__() + self.dolls = dolls + for doll in dolls: + setattr(self, f"mrl_{doll}", nn.Linear(doll, features)) + + def forward(self, x): + "x: (batch, features)" + logits = [getattr(self, f"mrl_{doll}")(x[:, :doll]) for doll in self.dolls] + return logits + + +class MRLLoss(nn.Module): + def __init__(self, weights) -> None: + super().__init__() + self.weights = weights + self.criterion = nn.CosineSimilarity(dim=1, eps=1e-6) + + def forward(self, representations, targets): + """ + representations: [(batch, features), ...] + targets: (batch, features) + """ + losses = [ + self.weights[i] * (1 - self.criterion(rep, targets)).mean() + for i, rep in enumerate(representations) + ] + return sum(losses) / len(losses) From 94f2b758f6ecfb72a86adc7a952a3157f6c5ebf5 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 20:45:21 +0530 Subject: [PATCH 19/83] SAM as teacher, MRL, split code into modules --- configs/config.yaml | 6 +- src/model.py | 147 ++++++-------------------------------------- src/module.py | 116 ++++++++++++++++++++++++++++++++++ trainer.py | 2 +- 4 files changed, 140 insertions(+), 131 deletions(-) create mode 100644 src/module.py diff --git a/configs/config.yaml b/configs/config.yaml index 217bf7fe..0c97f1b5 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -20,7 +20,9 @@ model: patch_size: 8 shuffle: True metadata_path: configs/metadata.yaml - teacher: vit_base_patch16_224.dino + teacher: samvit_base_patch16.sa1b + dolls: [16, 32, 64, 128, 256, 768] + doll_weights: [1, 1, 1, 1, 1, 1] lr: 1e-5 wd: 0.05 b1: 0.9 @@ -32,7 +34,7 @@ trainer: devices: auto num_nodes: 1 precision: bf16-mixed - log_every_n_steps: 10 + log_every_n_steps: 1 max_epochs: 200 accumulate_grad_batches: 1 default_root_dir: s3://clay-model-ckpt/v1.0.0/ diff --git a/src/model.py b/src/model.py index ee211b97..3f2b18c4 100644 --- a/src/model.py +++ b/src/model.py @@ -1,19 +1,15 @@ import math import os -from typing import Literal -import lightning as L import timm import torch import torch.nn.functional as F -import yaml -from box import Box from einops import rearrange, reduce, repeat from torch import nn -from torchvision.transforms import v2 -from vit_pytorch.simple_vit import Transformer +from src.backbone import Transformer from src.factory import DynamicEmbedding +from src.mrl import MRL, MRLLoss from src.utils import posemb_sincos_2d_with_gsd torch.set_float32_matmul_precision("medium") @@ -53,6 +49,7 @@ def __init__( # noqa: PLR0913 heads=heads, dim_head=dim_head, mlp_dim=int(dim * mlp_ratio), + fused_attn=True, ) def to_patch_embed(self, cube, waves): @@ -239,6 +236,7 @@ def __init__( # noqa: PLR0913 heads=heads, dim_head=dim_head, mlp_dim=int(dim * mlp_ratio), + fused_attn=True, ) self.embed_to_pixels = DynamicEmbedding( wave_dim=128, @@ -364,6 +362,8 @@ def __init__( # noqa: PLR0913 shuffle, metadata, teacher, + dolls, + doll_weights, # ENCODER dim, depth, @@ -385,11 +385,8 @@ def __init__( # noqa: PLR0913 self.shuffle = shuffle self.metadata = metadata self.teacher = timm.create_model(teacher, pretrained=True, num_classes=0) - self.teacher_chip_size = 224 - self.teacher_resize = v2.Resize( - size=(self.teacher_chip_size, self.teacher_chip_size) - ) - self.proj = nn.Linear(dim, self.teacher.num_features) + self.mrl = MRL(features=self.teacher.num_features, dolls=dolls) + self.mrl_loss = MRLLoss(weights=doll_weights) self.encoder = Encoder( mask_ratio=mask_ratio, @@ -418,6 +415,7 @@ def __init__( # noqa: PLR0913 def freeze_teacher(self): for param in self.teacher.parameters(): param.requires_grad = False + self.teacher.eval() def per_pixel_loss(self, cube, pixels, masked_matrix): """ @@ -487,13 +485,14 @@ def forward(self, datacube): waves, ) # [B L (C P P)] - # LOSS + # MAE reconstruction_loss = self.per_pixel_loss( datacube["pixels"], pixels, masked_matrix ) - # TEACHER - encoder_output = self.proj(encoded_unmasked_patches[:, 0, :]) # [B D'] + # MRL + representations = self.mrl(encoded_unmasked_patches[:, 0, :]) # [(B D') ...] + with torch.no_grad(): if platform == "sentinel-1-rtc": r = datacube["pixels"][:, 0, :, :] @@ -504,13 +503,9 @@ def forward(self, datacube): # Read RGB bands from the sensor to feed the teacher model indices = self.metadata[platform].rgb_indices rgb = datacube["pixels"][:, indices, :, :] - rgb = self.teacher_resize(rgb) - teacher_output = self.teacher(rgb) + target = self.teacher(rgb) - representation_loss = -( - F.cosine_similarity(encoder_output, teacher_output).mean() - - 1.0 # change range from [-1, 1] to [-2, 0] - ) # negative cosine similarity, [0, 2] -> 0 is similar & 2 is opposite + representation_loss = self.mrl_loss(representations, target) loss = 0.90 * reconstruction_loss + 0.10 * representation_loss return (loss, reconstruction_loss, representation_loss) @@ -563,11 +558,11 @@ def clay_mae_base(**kwargs): "dim_head": 64, "mlp_ratio": 4, # DECODER - "decoder_dim": 512, - "decoder_depth": 6, - "decoder_heads": 6, + "decoder_dim": 256, + "decoder_depth": 2, + "decoder_heads": 2, "decoder_dim_head": 64, - "decoder_mlp_ratio": 4, + "decoder_mlp_ratio": 2, } args.update(kwargs) return ClayMAE(**args) @@ -590,107 +585,3 @@ def clay_mae_large(**kwargs): } args.update(kwargs) return ClayMAE(**args) - - -class ClayMAEModule(L.LightningModule): - def __init__( # noqa: PLR0913 - self, - model_size="base", - mask_ratio=0.75, - norm_pix_loss=False, - patch_size=16, - shuffle=False, - metadata_path="configs/metadata.yaml", - teacher="vit_base_patch16_224.dino", - lr=1e-4, - wd=0.05, - b1=0.9, - b2=0.95, - embeddings_level: Literal["mean", "patch", "group"] = "mean", - ): - super().__init__() - self.save_hyperparameters(logger=True) - self.metadata = Box(yaml.safe_load(open(metadata_path))) - model_map = { - "tiny": clay_mae_tiny, - "small": clay_mae_small, - "base": clay_mae_base, - "large": clay_mae_large, - } - if model_size in model_map: - model_args = { - "mask_ratio": mask_ratio, - "patch_size": patch_size, - "norm_pix_loss": norm_pix_loss, - "shuffle": shuffle, - "metadata": self.metadata, - "teacher": teacher, - } - self.model = model_map[model_size](**model_args) - else: - raise ValueError( - f"Invalid model size {model_size}. Expected one of {model_map.keys()}" - ) - - def on_train_epoch_start(self): - self.model.teacher.eval() - - def forward(self, datacube: dict[str, torch.Tensor]): - return self.model(datacube) - - def configure_optimizers(self): - optimizer = torch.optim.AdamW( - self.parameters(), - lr=self.hparams.lr, - weight_decay=self.hparams.wd, - 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 * 100, last_epoch=-1 - ) - - return { - "optimizer": optimizer, - "lr_scheduler": { - "scheduler": scheduler, - "interval": "step", - }, - } - - def shared_step(self, batch: dict[str, torch.Tensor], batch_idx: int, phase: str): - datacube = batch - loss, reconstruction_loss, representation_loss = self(datacube) - self.log( - name=f"{phase}/loss", - value=loss, - on_step=True, - on_epoch=True, - prog_bar=True, - logger=True, - sync_dist=True, - ) - self.log( - name=f"{phase}/rec_loss", - value=reconstruction_loss, - on_step=True, - on_epoch=True, - prog_bar=True, - logger=True, - sync_dist=True, - ) - self.log( - name=f"{phase}/rep_loss", - value=representation_loss, - on_step=True, - on_epoch=True, - prog_bar=True, - logger=True, - sync_dist=True, - ) - return loss - - def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int): - return self.shared_step(batch, batch_idx, phase="train") - - def validation_step(self, batch: dict[str, torch.Tensor], batch_idx: int): - return self.shared_step(batch, batch_idx, phase="val") diff --git a/src/module.py b/src/module.py new file mode 100644 index 00000000..d352365b --- /dev/null +++ b/src/module.py @@ -0,0 +1,116 @@ +from typing import Literal + +import lightning as L +import torch +import yaml +from box import Box + +from src.model import clay_mae_base, clay_mae_large, clay_mae_small, clay_mae_tiny + + +class ClayMAEModule(L.LightningModule): + def __init__( # noqa: PLR0913 + self, + model_size="base", + mask_ratio=0.75, + norm_pix_loss=False, + patch_size=8, + shuffle=False, + metadata_path="configs/metadata.yaml", + teacher="samvit_base_patch16.sa1b", + dolls=[16, 32, 64, 128, 256, 768], + doll_weights=[1, 1, 1, 1, 1, 1], + lr=1e-5, + wd=0.05, + b1=0.9, + b2=0.95, + embeddings_level: Literal["mean", "patch", "group"] = "mean", + ): + super().__init__() + self.save_hyperparameters(logger=True) + self.metadata = Box(yaml.safe_load(open(metadata_path))) + model_map = { + "tiny": clay_mae_tiny, + "small": clay_mae_small, + "base": clay_mae_base, + "large": clay_mae_large, + } + if model_size in model_map: + model_args = { + "mask_ratio": mask_ratio, + "patch_size": patch_size, + "norm_pix_loss": norm_pix_loss, + "shuffle": shuffle, + "metadata": self.metadata, + "teacher": teacher, + "dolls": dolls, + "doll_weights": doll_weights, + } + self.model = model_map[model_size](**model_args) + else: + raise ValueError( + f"Invalid model size {model_size}. Expected one of {model_map.keys()}" + ) + + def on_train_epoch_start(self): + self.model.teacher.eval() + + def forward(self, datacube: dict[str, torch.Tensor]): + return self.model(datacube) + + def configure_optimizers(self): + optimizer = torch.optim.AdamW( + self.parameters(), + lr=self.hparams.lr, + weight_decay=self.hparams.wd, + 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 * 100, last_epoch=-1 + ) + + return { + "optimizer": optimizer, + "lr_scheduler": { + "scheduler": scheduler, + "interval": "step", + }, + } + + def shared_step(self, batch: dict[str, torch.Tensor], batch_idx: int, phase: str): + platform = batch["platform"][0] + loss, reconstruction_loss, representation_loss = self(batch) + + losses = { + "loss": loss, + "rec_loss": reconstruction_loss, + "rep_loss": representation_loss, + } + + for loss_name, loss_value in losses.items(): + self.log( + name=f"{phase}/{loss_name}", + value=loss_value, + on_step=True, + on_epoch=True, + prog_bar=True, + logger=True, + sync_dist=True, + ) + self.log( + name=f"{phase}_{platform}/{loss_name}", + value=loss_value, + on_step=True, + on_epoch=True, + prog_bar=True, + logger=True, + sync_dist=True, + ) + + return loss + + def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int): + return self.shared_step(batch, batch_idx, phase="train") + + def validation_step(self, batch: dict[str, torch.Tensor], batch_idx: int): + return self.shared_step(batch, batch_idx, phase="val") diff --git a/trainer.py b/trainer.py index 986574e8..77252a14 100644 --- a/trainer.py +++ b/trainer.py @@ -13,7 +13,7 @@ from lightning.pytorch.cli import LightningCLI from src.datamodule import ClayDataModule # noqa: F401 -from src.model import ClayMAEModule # noqa: F401 +from src.module import ClayMAEModule # noqa: F401 # %% From 08df662c9f284a0ce2d9a224ba6637d5a7d2843a Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 6 Aug 2024 20:46:44 +0530 Subject: [PATCH 20/83] Remove outdated README, all in docs --- src/README.md | 23 ----------------------- 1 file changed, 23 deletions(-) delete mode 100644 src/README.md diff --git a/src/README.md b/src/README.md deleted file mode 100644 index 7326daf7..00000000 --- a/src/README.md +++ /dev/null @@ -1,23 +0,0 @@ -# Clay Foundation Model Modules - -This folder contains several LightningDataModule, LightningModule and callback -classes. - -## DataModules (data pipeline) - -- datamodule.py - Data pipeline to read in Earth Observation chips from GeoTIFF files - -## LightningModule (model architecture) - -- model_clay.py - Clay Foundation Model architecture with spatiotemporal encoders -- model_vit.py - Vanilla Vision Transformer neural network model architecture - -## Callbacks (custom plugins) - -- callbacks_wandb.py - Log metrics and predictions to Weights and Biases while training. - -## References - -- https://lightning.ai/docs/pytorch/2.1.0/data/datamodule.html -- https://lightning.ai/docs/pytorch/2.1.0/common/lightning_module.html -- https://lightning.ai/docs/pytorch/2.1.0/extensions/callbacks.html From c6df3e66360296152962d9c193b13f27dc7f7bfc Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 7 Aug 2024 10:46:47 +0000 Subject: [PATCH 21/83] Fix trainer --- trainer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/trainer.py b/trainer.py index 77252a14..cce9ecac 100644 --- a/trainer.py +++ b/trainer.py @@ -21,7 +21,11 @@ def cli_main(): """ Command-line inteface to run ClayMAE with ClayDataModule. """ - cli = LightningCLI(save_config_kwargs={"overwrite": True}) + cli = LightningCLI( + ClayMAEModule, + ClayDataModule, + save_config_kwargs={"overwrite": True} + ) return cli From c6151c6f273419e5d6efbe9b882487cc8397f3eb Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 7 Aug 2024 10:47:37 +0000 Subject: [PATCH 22/83] Fix mean, std for s1 --- configs/metadata.yaml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/configs/metadata.yaml b/configs/metadata.yaml index 7e3d71af..5371487a 100644 --- a/configs/metadata.yaml +++ b/configs/metadata.yaml @@ -176,11 +176,11 @@ sentinel-1-rtc: gsd: 10 bands: mean: - vv: 0.123273 - vh: 0.027337 + vv: 0.123 + vh: 0.027 std: - vv: 1.492154 - vh: 0.122182 + vv: 0.689 + vh: 0.061 wavelength: vv: 3.5 vh: 4.0 From a8cee0f2c304fa63a509d4f31f3acfb2f72c5b28 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 7 Aug 2024 10:48:41 +0000 Subject: [PATCH 23/83] update config for 1 node run --- configs/config.yaml | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 0c97f1b5..446d8f3c 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,22 +1,22 @@ -# lightning.pytorch==2.1.2 -seed_everything: 42 +seed_everything: 108 data: - data_dir: data + data_dir: /opt/dlami/nvme/pretrain/ size: 224 metadata_path: configs/metadata.yaml platforms: - landsat-c2l1 - landsat-c2l2-sr - linz + - modis - naip - sentinel-1-rtc - sentinel-2-l2a - batch_size: 8 - num_workers: 8 + batch_size: 1 + num_workers: 4 model: model_size: base mask_ratio: 0.75 - norm_pix_loss: True + norm_pix_loss: False patch_size: 8 shuffle: True metadata_path: configs/metadata.yaml @@ -37,7 +37,7 @@ trainer: log_every_n_steps: 1 max_epochs: 200 accumulate_grad_batches: 1 - default_root_dir: s3://clay-model-ckpt/v1.0.0/ + default_root_dir: checkpoints/v1.0.1/ fast_dev_run: False num_sanity_val_steps: 0 use_distributed_sampler: False @@ -50,19 +50,19 @@ trainer: callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - dirpath: s3://clay-model-ckpt/v1.0.0/ + dirpath: checkpoints/v1.0.1/ auto_insert_metric_name: False - filename: mae_v1.0.0_epoch-{epoch:02d}_val-loss-{val/loss:.4f} + filename: mae_v1.0.1_epoch-{epoch:02d}_val-loss-{val/loss:.4f} monitor: val/loss mode: min save_last: True save_top_k: 2 save_weights_only: False verbose: True - - class_path: lightning.pytorch.callbacks.LearningRateMonitor - init_args: - logging_interval: step - - class_path: src.callbacks_wandb.LogIntermediatePredictions + # - class_path: lightning.pytorch.callbacks.LearningRateMonitor + # init_args: + # logging_interval: step + # - class_path: src.callbacks_wandb.LogIntermediatePredictions plugins: - class_path: lightning.pytorch.plugins.io.AsyncCheckpointIO ckpt_path: null From 4dceba35cfdea17ec090431f91594020e044a2c5 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 8 Aug 2024 06:15:44 +0000 Subject: [PATCH 24/83] Modify Sentinel 1 from raw pixels to dB scale --- configs/config.yaml | 8 ++++---- configs/metadata.yaml | 8 ++++---- copy_data.sh | 18 ++++++++++++++++++ src/callbacks_wandb.py | 10 ++++++---- src/datamodule.py | 9 ++++++++- src/model.py | 4 ++-- 6 files changed, 42 insertions(+), 15 deletions(-) create mode 100644 copy_data.sh diff --git a/configs/config.yaml b/configs/config.yaml index 446d8f3c..3cdc9678 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -59,10 +59,10 @@ trainer: save_top_k: 2 save_weights_only: False verbose: True - # - class_path: lightning.pytorch.callbacks.LearningRateMonitor - # init_args: - # logging_interval: step - # - class_path: src.callbacks_wandb.LogIntermediatePredictions + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: src.callbacks_wandb.LogIntermediatePredictions plugins: - class_path: lightning.pytorch.plugins.io.AsyncCheckpointIO ckpt_path: null diff --git a/configs/metadata.yaml b/configs/metadata.yaml index 5371487a..193bc438 100644 --- a/configs/metadata.yaml +++ b/configs/metadata.yaml @@ -176,11 +176,11 @@ sentinel-1-rtc: gsd: 10 bands: mean: - vv: 0.123 - vh: 0.027 + vv: -12.113 + vh: -18.673 std: - vv: 0.689 - vh: 0.061 + vv: 8.314 + vh: 8.017 wavelength: vv: 3.5 vh: 4.0 diff --git a/copy_data.sh b/copy_data.sh new file mode 100644 index 00000000..43c63516 --- /dev/null +++ b/copy_data.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +# Define source and destination directories +src="/fsx" +dest="data/pretrain" + +# Create the destination directory if it doesn't exist +mkdir -p "$dest" + +# Find all directories in the source directory +find "$src" -type d -print0 | while IFS= read -r -d '' dir; do + # Create corresponding directory in the destination + newdir="$dest${dir#$src}" + mkdir -p "$newdir" + + # Copy the first 100 files from the source directory to the new directory + find "$dir" -maxdepth 1 -type f -print0 | head -z -n 100 | xargs -0 -I{} cp {} "$newdir" +done diff --git a/src/callbacks_wandb.py b/src/callbacks_wandb.py index 0f4d4a10..15ec437c 100644 --- a/src/callbacks_wandb.py +++ b/src/callbacks_wandb.py @@ -247,6 +247,8 @@ def on_validation_end( ) assert pixels.shape == batch["pixels"].shape + batch["pixels"] = batch["pixels"].detach().cpu().numpy() + pixels = pixels.detach().cpu().numpy() n_rows = 4 # 2 for actual and 2 for predicted n_cols = 8 @@ -256,13 +258,13 @@ def on_validation_end( for j in range(n_cols): # Plot actual images in rows 0 and 2 axs[0, j].imshow( - batch["pixels"][j][0].detach().cpu().numpy(), cmap="viridis" + batch["pixels"][j][0], cmap="viridis" ) axs[0, j].set_title(f"Actual {j}") axs[0, j].axis("off") axs[2, j].imshow( - batch["pixels"][j + n_cols][0].detach().cpu().numpy(), + batch["pixels"][j + n_cols][0], cmap="viridis", ) axs[2, j].set_title(f"Actual {j+n_cols}") @@ -270,13 +272,13 @@ def on_validation_end( # Plot predicted images in rows 1 and 3 axs[1, j].imshow( - pixels[j][0].detach().cpu().numpy(), cmap="viridis" + pixels[j][0], cmap="viridis" ) axs[1, j].set_title(f"Pred {j}") axs[1, j].axis("off") axs[3, j].imshow( - pixels[j + n_cols][0].detach().cpu().numpy(), cmap="viridis" + pixels[j + n_cols][0], cmap="viridis" ) axs[3, j].set_title(f"Pred {j+n_cols}") axs[3, j].axis("off") diff --git a/src/datamodule.py b/src/datamodule.py index 7cfe579f..75e3e499 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -54,8 +54,15 @@ def __len__(self): def __getitem__(self, idx): chip_path = self.chips_path[idx] with np.load(chip_path, allow_pickle=False) as chip: - pixels = torch.from_numpy(chip["pixels"].astype(np.float32)) platform = chip_path.parent.name + if platform == "sentinel-1-rtc": + pixels = chip["pixels"].astype(np.float32) + pixels[pixels <= 0] = 1e-10 # replace corrupted pixels in sentinel-1-rtc with small value + pixels = 10 * np.log10(pixels) # convert to dB scale, more interpretable pixels + else: + pixels = chip["pixels"].astype(np.float32) + + pixels = torch.from_numpy(pixels) pixels = self.transforms[platform](pixels) time_tensor = torch.tensor( diff --git a/src/model.py b/src/model.py index 3f2b18c4..d9a94e8d 100644 --- a/src/model.py +++ b/src/model.py @@ -497,7 +497,7 @@ def forward(self, datacube): if platform == "sentinel-1-rtc": r = datacube["pixels"][:, 0, :, :] g = datacube["pixels"][:, 1, :, :] - b = r - g + b = (r + g)/2 rgb = torch.stack((r, g, b), dim=1) else: # Read RGB bands from the sensor to feed the teacher model @@ -507,7 +507,7 @@ def forward(self, datacube): representation_loss = self.mrl_loss(representations, target) - loss = 0.90 * reconstruction_loss + 0.10 * representation_loss + loss = 0.95 * reconstruction_loss + 0.05 * representation_loss return (loss, reconstruction_loss, representation_loss) From fe0167362d4daada04e76ba66fce6fee04a020bc Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Thu, 8 Aug 2024 17:16:47 +0530 Subject: [PATCH 25/83] Cluster template for multi-node training --- cluster/ml-cluster.yaml.template | 61 ++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 cluster/ml-cluster.yaml.template diff --git a/cluster/ml-cluster.yaml.template b/cluster/ml-cluster.yaml.template new file mode 100644 index 00000000..03a12836 --- /dev/null +++ b/cluster/ml-cluster.yaml.template @@ -0,0 +1,61 @@ +Region: us-east-2 + +# DL AMI +Image: + Os: ubuntu2004 + CustomAmi: + +# FSx LUSTRE SHARED STORAGE +SharedStorage: + - MountDir: /fsx + Name: fsx + StorageType: FsxLustre + FsxLustreSettings: + FileSystemId: + +# HEAD NODE +HeadNode: + InstanceType: c5.12xlarge + Networking: + SubnetId: + SecurityGroups: + - # EFA enabled SG + Ssh: + KeyName: + LocalStorage: + RootVolume: + Size: 200 + Iam: + S3Access: + - BucketName: + EnableWriteAccess: false + - BucketName: + EnableWriteAccess: true + + +# SCHEDULER +Scheduling: + Scheduler: slurm + SlurmQueues: + - Name: gpu-queue + ComputeResources: + - Name: + Instances: + - InstanceType: + MinCount: 0 + MaxCount: 8 + Efa: + Enabled: true + Networking: + SubnetIds: + - + SecurityGroups: + - # EFA enabled SG + PlacementGroup: + Enabled: true + Iam: + S3Access: + - BucketName: + EnableWriteAccess: false + - BucketName: + EnableWriteAccess: true From 940fbab3a4c26876823de45f615f675434db7b08 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Mon, 12 Aug 2024 12:53:22 +0530 Subject: [PATCH 26/83] Fix distributed DataLoader, add new env & slurm script --- configs/config.yaml | 20 +++++---- src/datamodule.py | 101 +++++++++++++++++++++++++++++++++++++----- src/model.py | 10 ++--- src/module.py | 3 +- train_clay_v2.sh | 67 ++++++++++++++++++++++++++++ train_environment.yml | 21 +++++++++ 6 files changed, 195 insertions(+), 27 deletions(-) create mode 100644 train_clay_v2.sh create mode 100644 train_environment.yml diff --git a/configs/config.yaml b/configs/config.yaml index 3cdc9678..5736666c 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,6 +1,6 @@ seed_everything: 108 data: - data_dir: /opt/dlami/nvme/pretrain/ + data_dir: /fsx size: 224 metadata_path: configs/metadata.yaml platforms: @@ -12,7 +12,7 @@ data: - sentinel-1-rtc - sentinel-2-l2a batch_size: 1 - num_workers: 4 + num_workers: 12 model: model_size: base mask_ratio: 0.75 @@ -23,24 +23,26 @@ model: teacher: samvit_base_patch16.sa1b dolls: [16, 32, 64, 128, 256, 768] doll_weights: [1, 1, 1, 1, 1, 1] - lr: 1e-5 + lr: 5e-6 wd: 0.05 b1: 0.9 b2: 0.95 embeddings_level: mean trainer: - accelerator: auto + accelerator: gpu strategy: ddp - devices: auto - num_nodes: 1 + devices: 4 + num_nodes: 2 precision: bf16-mixed log_every_n_steps: 1 max_epochs: 200 accumulate_grad_batches: 1 - default_root_dir: checkpoints/v1.0.1/ + default_root_dir: checkpoints/v1.0.4/ fast_dev_run: False num_sanity_val_steps: 0 use_distributed_sampler: False + limit_train_batches: 0.99 + limit_val_batches: 0.99 logger: - class_path: lightning.pytorch.loggers.WandbLogger init_args: @@ -50,9 +52,9 @@ trainer: callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - dirpath: checkpoints/v1.0.1/ + dirpath: checkpoints/v1.0.4/ auto_insert_metric_name: False - filename: mae_v1.0.1_epoch-{epoch:02d}_val-loss-{val/loss:.4f} + filename: mae_v1.0.4_epoch-{epoch:02d}_val-loss-{val/loss:.4f} monitor: val/loss mode: min save_last: True diff --git a/src/datamodule.py b/src/datamodule.py index 75e3e499..0670f2d1 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -3,6 +3,7 @@ rasterio. """ +import math import random from collections import defaultdict from pathlib import Path @@ -11,7 +12,8 @@ import lightning as L import numpy as np import torch -import torchdata + +# import torchdata import yaml from box import Box from einops import rearrange @@ -57,8 +59,12 @@ def __getitem__(self, idx): platform = chip_path.parent.name if platform == "sentinel-1-rtc": pixels = chip["pixels"].astype(np.float32) - pixels[pixels <= 0] = 1e-10 # replace corrupted pixels in sentinel-1-rtc with small value - pixels = 10 * np.log10(pixels) # convert to dB scale, more interpretable pixels + pixels[pixels <= 0] = ( + 1e-10 # replace corrupted pixels in sentinel-1-rtc with small value + ) + pixels = 10 * np.log10( + pixels + ) # convert to dB scale, more interpretable pixels else: pixels = chip["pixels"].astype(np.float32) @@ -121,6 +127,77 @@ def __len__(self): return len(self.dataset.chips_path) // self.batch_size +class ClayDistributedSampler(Sampler): + def __init__( # noqa: PLR0913 + self, + dataset, + platforms, + batch_size, + num_replicas=None, + rank=None, + shuffle=True, + ): + self.dataset = dataset + self.platforms = platforms + self.batch_size = batch_size + self.num_replicas = ( + num_replicas + if num_replicas is not None + else torch.distributed.get_world_size() + ) + self.rank = rank if rank is not None else torch.distributed.get_rank() + self.shuffle = shuffle + self.epoch = 0 + + self.platform_indices = {platform: [] for platform in platforms} + for idx, chip_path in enumerate(self.dataset.chips_path): + platform = chip_path.parent.name + self.platform_indices[platform].append(idx) + + self.max_len = max(len(indices) for indices in self.platform_indices.values()) + self.adjusted_indices = {} + # Normalize the length of indices for each platform by replicating the indices + # to match the max_len + for platform, indices in self.platform_indices.items(): + if len(indices) < self.max_len: + extended_indices = np.tile(indices, (self.max_len // len(indices) + 1))[ + : self.max_len + ] + self.adjusted_indices[platform] = extended_indices + else: + self.adjusted_indices[platform] = indices + + self.num_samples = math.ceil( + ((self.max_len * len(self.platforms)) - self.num_replicas) + / self.num_replicas + ) + self.total_size = self.num_samples * self.num_replicas + self.num_samples_per_platform = self.max_len // self.num_replicas + + def __iter__(self): + rng = np.random.default_rng(self.epoch) + platform_batches = {} + for platform, indices in self.adjusted_indices.items(): + if self.shuffle: + rng.shuffle(indices) + # Distribute the indices to each process + start_idx = self.rank * self.num_samples_per_platform + end_idx = start_idx + self.num_samples_per_platform + platform_batches[platform] = indices[start_idx:end_idx] + + for i in range(0, self.num_samples_per_platform, self.batch_size): + for platform in self.platforms: + batch = platform_batches[platform][i : i + self.batch_size] + if len(batch) == self.batch_size: + yield batch + + def __len__(self) -> int: + return self.num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def batch_collate(batch): """Collate function for DataLoader. @@ -171,13 +248,13 @@ def __init__( # noqa: PLR0913 def setup(self, stage: Literal["fit", "predict"] | None = None) -> None: # Get list of GeoTIFF filepaths from s3 bucket or data/ folder - if self.data_dir.startswith("s3://"): - dp = torchdata.datapipes.iter.IterableWrapper(iterable=[self.data_dir]) - chips_path = list(dp.list_files_by_s3(masks="*.npz")) - else: # if self.data_dir is a local data path - chips_path = sorted(list(Path(self.data_dir).glob("**/*.npz"))) - chips_platform = [chip.parent.parent.name for chip in chips_path] - # chips_platform = [chip.parent.parent.name for chip in chips_path] + # if self.data_dir.startswith("s3://"): + # dp = torchdata.datapipes.iter.IterableWrapper(iterable=[self.data_dir]) + # chips_path = list(dp.list_files_by_s3(masks="*.npz")) + # else: # if self.data_dir is a local data path + chips_path = sorted(list(Path(self.data_dir).glob("**/*.npz"))) + chips_platform = [chip.parent.parent.name for chip in chips_path] + # chips_platform = [chip.parent.parent.name for chip in chips_path] print(f"Total number of chips: {len(chips_path)}") if stage == "fit": @@ -194,7 +271,7 @@ def setup(self, stage: Literal["fit", "predict"] | None = None) -> None: platforms=self.platforms, metadata=self.metadata, ) - self.trn_sampler = ClaySampler( + self.trn_sampler = ClayDistributedSampler( dataset=self.trn_ds, platforms=self.platforms, batch_size=self.batch_size, @@ -205,7 +282,7 @@ def setup(self, stage: Literal["fit", "predict"] | None = None) -> None: platforms=self.platforms, metadata=self.metadata, ) - self.val_sampler = ClaySampler( + self.val_sampler = ClayDistributedSampler( dataset=self.val_ds, platforms=self.platforms, batch_size=self.batch_size, diff --git a/src/model.py b/src/model.py index d9a94e8d..711bc8f7 100644 --- a/src/model.py +++ b/src/model.py @@ -497,7 +497,7 @@ def forward(self, datacube): if platform == "sentinel-1-rtc": r = datacube["pixels"][:, 0, :, :] g = datacube["pixels"][:, 1, :, :] - b = (r + g)/2 + b = (r + g) / 2 rgb = torch.stack((r, g, b), dim=1) else: # Read RGB bands from the sensor to feed the teacher model @@ -558,11 +558,11 @@ def clay_mae_base(**kwargs): "dim_head": 64, "mlp_ratio": 4, # DECODER - "decoder_dim": 256, - "decoder_depth": 2, - "decoder_heads": 2, + "decoder_dim": 512, + "decoder_depth": 4, + "decoder_heads": 4, "decoder_dim_head": 64, - "decoder_mlp_ratio": 2, + "decoder_mlp_ratio": 4, } args.update(kwargs) return ClayMAE(**args) diff --git a/src/module.py b/src/module.py index d352365b..f247d51a 100644 --- a/src/module.py +++ b/src/module.py @@ -64,9 +64,10 @@ def configure_optimizers(self): lr=self.hparams.lr, weight_decay=self.hparams.wd, betas=(self.hparams.b1, self.hparams.b2), + fused=True, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( - optimizer, T_0=1000, T_mult=2, eta_min=self.hparams.lr * 100, last_epoch=-1 + optimizer, T_0=2000, T_mult=1, eta_min=self.hparams.lr * 100, last_epoch=-1 ) return { diff --git a/train_clay_v2.sh b/train_clay_v2.sh new file mode 100644 index 00000000..cb5abbb3 --- /dev/null +++ b/train_clay_v2.sh @@ -0,0 +1,67 @@ +#!/bin/bash + +#SBATCH --job-name=clay-laucher +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=4 # EDIT if it's not 8-gpus per node +#SBATCH --cpus-per-task=12 # EDIT this to how many cpu cores the node has divided by num of gpus +#SBATCH --gres=gpu:4 # EDIT this if it's not 8-gpus per node +#SBATCH --time=0-00:00:00 # EDIT the desired runtime +#SBATCH --exclusive +#SBATCH --partition=gpu # EDIT to the desired partition name +#SBATCH --output=%x-%j-%N.out + +echo "START TIME: $(date)" + +# auto-fail on any errors in this script +set -eo pipefail + +# logging script's variables/commands for future debug needs +set -x + +# EDIT the conda evn and any startup scripts +# source /path/to/start-xxx-user # if you have something to preload before the job +# Load any required modules (environments, libraries etc.) +eval "$(conda 'shell.bash' 'hook' 2> /dev/null)" + +# initialize conda +conda activate /home/ubuntu/claymodel # if you have conda env to activate + +LOG_PATH="main_log.txt" + +# PTL doesn't need a special launcher +LAUNCHER="python -u" + +# EDIT the path+name of the python script and whatever args it needs +PROGRAM="trainer.py fit --config configs/config.yaml" + +export CMD="$LAUNCHER $PROGRAM" + +echo $CMD + +# EDIT if you want to redirect /tmp to /scratch (some local SSD path) since /tmp is tiny on compute nodes +# export TMPDIR=/scratch + +# EDIT: useful for debug if needed +# +# to debug NCCL issues +# export NCCL_DEBUG=INFO +# +# to unravel async errors w/o the correct traceback - potentially makes everything very slower +# export CUDA_LAUNCH_BLOCKING=1 +# +# to force crashing on nccl issues like hanging broadcast +# export NCCL_ASYNC_ERROR_HANDLING=1 + +# srun error handling: +# --wait=60: wait 60 sec after the first task terminates before terminating all remaining tasks +# --kill-on-bad-exit=1: terminate a step if any task exits with a non-zero exit code +SRUN_ARGS=" \ + --wait=60 \ + --kill-on-bad-exit=1 \ + --jobid $SLURM_JOB_ID \ + " + +# bash -c is needed for the delayed interpolation of env vars to work +srun $SRUN_ARGS bash -c "$CMD" 2>&1 | tee -a $LOG_PATH + +echo "END TIME: $(date)" diff --git a/train_environment.yml b/train_environment.yml new file mode 100644 index 00000000..433e5b6c --- /dev/null +++ b/train_environment.yml @@ -0,0 +1,21 @@ +name: claymodel +channels: + - pytorch + - conda-forge +dependencies: + - python=3.11 + - pip + - pip: + - einops~=0.7.0 + - geopandas + - jsonargparse[signatures]>=4.27.7 + - lightning + - matplotlib + - python-box + - torch + - scikit-image + - scikit-learn + - timm + - torchvision + - vit-pytorch + - wandb From f12e22151bc4b04098e075fa66f2396580587a49 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Mon, 12 Aug 2024 12:56:07 +0530 Subject: [PATCH 27/83] Fix docs --- src/backbone.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backbone.py b/src/backbone.py index 16fd1fd5..a6e2ebb2 100644 --- a/src/backbone.py +++ b/src/backbone.py @@ -1,4 +1,4 @@ -"""Code for Transformer from Phil Wangs library +"""Code for Transformer from Phil Wangs vit-pytorch library. Repository: https://github.com/lucidrains/vit-pytorch """ From ea4136bcf675aa5872f3898ff043141e3836176d Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Tue, 13 Aug 2024 12:06:37 +0530 Subject: [PATCH 28/83] Scale down recontruction loss for MODIS, change alpha to 0.9 for rec/rep split --- src/model.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/model.py b/src/model.py index 711bc8f7..555f914b 100644 --- a/src/model.py +++ b/src/model.py @@ -489,6 +489,10 @@ def forward(self, datacube): reconstruction_loss = self.per_pixel_loss( datacube["pixels"], pixels, masked_matrix ) + # MODIS has a 10x reconstruction loss compared to all the other sensors, + # so we need to scale it down to improve the learning capability. + if platform == "modis": + reconstruction_loss /= 10 # MRL representations = self.mrl(encoded_unmasked_patches[:, 0, :]) # [(B D') ...] @@ -507,7 +511,7 @@ def forward(self, datacube): representation_loss = self.mrl_loss(representations, target) - loss = 0.95 * reconstruction_loss + 0.05 * representation_loss + loss = 0.9 * reconstruction_loss + 0.1 * representation_loss return (loss, reconstruction_loss, representation_loss) From 224c547d6bab1c2046380435d1eb859604a873d6 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 21 Aug 2024 11:49:43 +0530 Subject: [PATCH 29/83] Use groups in wandb --- configs/config.yaml | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 5736666c..5082ff9f 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -37,7 +37,7 @@ trainer: log_every_n_steps: 1 max_epochs: 200 accumulate_grad_batches: 1 - default_root_dir: checkpoints/v1.0.4/ + default_root_dir: checkpoints/v1.0.5/ fast_dev_run: False num_sanity_val_steps: 0 use_distributed_sampler: False @@ -48,13 +48,16 @@ trainer: init_args: entity: developmentseed project: clay + group: pixel-drop + # id: v8jh2pn9 + # resume: must log_model: false callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - dirpath: checkpoints/v1.0.4/ + dirpath: checkpoints/v1.0.5/ auto_insert_metric_name: False - filename: mae_v1.0.4_epoch-{epoch:02d}_val-loss-{val/loss:.4f} + filename: mae_v1.0.5_epoch-{epoch:02d}_val-loss-{val/loss:.4f} monitor: val/loss mode: min save_last: True From 79712a7d8d35c1313daead44b9ceb67411c15bc3 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 21 Aug 2024 11:53:33 +0530 Subject: [PATCH 30/83] Add random dropping for channels --- src/model.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/src/model.py b/src/model.py index 555f914b..630bb1ee 100644 --- a/src/model.py +++ b/src/model.py @@ -1,5 +1,6 @@ import math import os +import random import timm import torch @@ -457,6 +458,27 @@ def forward(self, datacube): waves = torch.tensor(list(self.metadata[platform].bands.wavelength.values())) gsd = torch.tensor(self.metadata[platform].gsd) + # Drop channels randomly + _pixels = datacube["pixels"].clone() + batch_size, channels, _, _ = _pixels.size() + + # Define probabilities for dropping channels + prob_drop_all = 0.10 # 10% probability to drop all channels + prob_drop_half = 0.20 # 20% probability to drop half the channels + + for i in range(batch_size): + if torch.any( + datacube["latlon"][i] != 0 + ): # Check if latlon is not all zeros + rand_val = random.random() + if rand_val < prob_drop_all: + _pixels[i, :, :, :] = 0 # Drop all channels + elif rand_val < prob_drop_all + prob_drop_half: + channel_indices = torch.randperm(channels)[ + : channels // 2 + ] # Get 50% of channel indices + _pixels[i, channel_indices, :, :] = 0 # Drop 50% of channels + # ENCODER ( encoded_unmasked_patches, # [B (1 + L):(1 - mask_ratio) D] @@ -465,7 +487,7 @@ def forward(self, datacube): masked_matrix, # [B L] ) = self.encoder( { - "pixels": datacube["pixels"], + "pixels": _pixels, "time": datacube["time"], "latlon": datacube["latlon"], "gsd": gsd, From 380e8145d4c9d26325f8e63f044b9cbe5b7abb6b Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 21 Aug 2024 18:59:36 +0530 Subject: [PATCH 31/83] Add script to check sanity of npz files --- utils/check_data_sanity.py | 60 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 utils/check_data_sanity.py diff --git a/utils/check_data_sanity.py b/utils/check_data_sanity.py new file mode 100644 index 00000000..a5276239 --- /dev/null +++ b/utils/check_data_sanity.py @@ -0,0 +1,60 @@ +import os +from concurrent.futures import ThreadPoolExecutor, as_completed + +import numpy as np + + +def check_and_delete_npz(file_path): + try: + # Attempt to load the .npz file using numpy + data = np.load(file_path) + + # Check if the 'pixel' key exists and has shape 128 in the 0th dimension + if "pixels" in data: + if data["pixels"].shape[0] != 128: # noqa: PLR2004 + os.remove(file_path) + return ( + None, + f"Invalid shape (not 128 in 0th dim): {file_path} - Deleted", + ) + else: + return f"Valid: {file_path}", None + else: + os.remove(file_path) + return None, f"'pixels' key missing: {file_path} - Deleted" + + except Exception as e: + os.remove(file_path) + return None, f"Invalid (Exception): {file_path} - {str(e)} - Deleted" + + +def process_directory_in_parallel(directory, max_workers=4): + invalid_files = [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [] + for root, dirs, files in os.walk(directory): + for file in files: + if file.endswith(".npz"): + file_path = os.path.join(root, file) + futures.append(executor.submit(check_and_delete_npz, file_path)) + + for future in as_completed(futures): + valid_msg, invalid_msg = future.result() + if valid_msg: + print(valid_msg) + if invalid_msg: + print(invalid_msg) + invalid_files.append(invalid_msg) + + return invalid_files + + +# Replace 'your_directory_path' with the path to the directory you want to check +invalid_files = process_directory_in_parallel("/fsx", max_workers=24) + +if invalid_files: + print("\nInvalid or corrupted .npz files found and deleted:") + for file in invalid_files: + print(file) +else: + print("\nAll .npz files are valid and meet the shape criteria for 'pixel' key.") From 4fe71384a19efaffd14550f196ab1956ea9563af Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Wed, 21 Aug 2024 19:08:47 +0530 Subject: [PATCH 32/83] Add script to split npz files of batch 128 to 32 --- utils/split_npz.py | 53 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 utils/split_npz.py diff --git a/utils/split_npz.py b/utils/split_npz.py new file mode 100644 index 00000000..cb2e87e5 --- /dev/null +++ b/utils/split_npz.py @@ -0,0 +1,53 @@ +import os +from concurrent.futures import ProcessPoolExecutor + +import numpy as np + + +def split_npz_file(file_path): + # Load the .npz file + with np.load(file_path) as data: + # Check if the file has the required batch size of 128 + if "pixels" in data and data["pixels"].shape[0] == 128: # noqa: PLR2004 + # Extract all arrays + keys = data.files + arrays = {key: data[key] for key in keys} + + # Determine the batch size and the number of splits + batch_size = 32 + num_splits = 4 # Since we want to split into 4 files, each with 32 samples + + # Split and save the smaller .npz files + for i in range(num_splits): + split_data = { + key: value[i * batch_size : (i + 1) * batch_size] + for key, value in arrays.items() + } + split_file_path = file_path.replace(".npz", f"_{i}.npz") + np.savez(split_file_path, **split_data) + print(f"Saved {split_file_path}") + + # Delete the original file + os.remove(file_path) + print(f"Deleted original file: {file_path}") + else: + print(f"Skipped {file_path}: Does not have a batch size of 128") + + +def process_directory(root_dir): + # Collect all .npz files + npz_files = [] + for dirpath, _, filenames in os.walk(root_dir): + for filename in filenames: + if filename.endswith(".npz"): + file_path = os.path.join(dirpath, filename) + npz_files.append(file_path) + + # Process files in parallel + with ProcessPoolExecutor() as executor: + executor.map(split_npz_file, npz_files) + + +# Example usage +root_dir = "/fsx" +process_directory(root_dir) From bb6d13c706ff371b8e8f89b54bcf6dcb3717fcb5 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Fri, 20 Sep 2024 23:09:24 +0530 Subject: [PATCH 33/83] Adapt script to large model size --- configs/config.yaml | 22 +++++++++++----------- src/datamodule.py | 2 +- src/model.py | 4 ++-- src/module.py | 2 +- train_clay_v2.sh | 6 +++--- 5 files changed, 18 insertions(+), 18 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 5082ff9f..1f9de9ee 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -1,7 +1,7 @@ seed_everything: 108 data: data_dir: /fsx - size: 224 + size: 256 metadata_path: configs/metadata.yaml platforms: - landsat-c2l1 @@ -14,15 +14,15 @@ data: batch_size: 1 num_workers: 12 model: - model_size: base + model_size: large mask_ratio: 0.75 norm_pix_loss: False patch_size: 8 shuffle: True metadata_path: configs/metadata.yaml teacher: samvit_base_patch16.sa1b - dolls: [16, 32, 64, 128, 256, 768] - doll_weights: [1, 1, 1, 1, 1, 1] + dolls: [16, 32, 64, 128, 256, 768, 1024] + doll_weights: [1, 1, 1, 1, 1, 1, 1] lr: 5e-6 wd: 0.05 b1: 0.9 @@ -31,13 +31,13 @@ model: trainer: accelerator: gpu strategy: ddp - devices: 4 - num_nodes: 2 + devices: 8 + num_nodes: 20 precision: bf16-mixed log_every_n_steps: 1 - max_epochs: 200 + max_epochs: 1000 accumulate_grad_batches: 1 - default_root_dir: checkpoints/v1.0.5/ + default_root_dir: checkpoints/v1.5.0/ fast_dev_run: False num_sanity_val_steps: 0 use_distributed_sampler: False @@ -48,16 +48,16 @@ trainer: init_args: entity: developmentseed project: clay - group: pixel-drop + group: v1.5 # id: v8jh2pn9 # resume: must log_model: false callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint init_args: - dirpath: checkpoints/v1.0.5/ + dirpath: checkpoints/v1.5.0/ auto_insert_metric_name: False - filename: mae_v1.0.5_epoch-{epoch:02d}_val-loss-{val/loss:.4f} + filename: mae_v1.5.0_epoch-{epoch:02d}_val-loss-{val/loss:.4f} monitor: val/loss mode: min save_last: True diff --git a/src/datamodule.py b/src/datamodule.py index 0670f2d1..ac7d7b5a 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -45,7 +45,7 @@ def create_transforms(self, mean, std): [ v2.RandomHorizontalFlip(p=0.5), v2.RandomVerticalFlip(p=0.5), - v2.RandomCrop(size=(self.size, self.size)), + # v2.RandomCrop(size=(self.size, self.size)), v2.Normalize(mean=mean, std=std), ] ) diff --git a/src/model.py b/src/model.py index 630bb1ee..797590d7 100644 --- a/src/model.py +++ b/src/model.py @@ -604,8 +604,8 @@ def clay_mae_large(**kwargs): "mlp_ratio": 4, # DECODER "decoder_dim": 512, - "decoder_depth": 8, - "decoder_heads": 8, + "decoder_depth": 4, + "decoder_heads": 4, "decoder_dim_head": 64, "decoder_mlp_ratio": 4, } diff --git a/src/module.py b/src/module.py index f247d51a..39b9df3a 100644 --- a/src/module.py +++ b/src/module.py @@ -67,7 +67,7 @@ def configure_optimizers(self): fused=True, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( - optimizer, T_0=2000, T_mult=1, eta_min=self.hparams.lr * 100, last_epoch=-1 + optimizer, T_0=5000, T_mult=1, eta_min=self.hparams.lr * 100, last_epoch=-1 ) return { diff --git a/train_clay_v2.sh b/train_clay_v2.sh index cb5abbb3..b39332e3 100644 --- a/train_clay_v2.sh +++ b/train_clay_v2.sh @@ -1,10 +1,10 @@ #!/bin/bash #SBATCH --job-name=clay-laucher -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=4 # EDIT if it's not 8-gpus per node +#SBATCH --nodes=20 +#SBATCH --ntasks-per-node=8 # EDIT if it's not 8-gpus per node #SBATCH --cpus-per-task=12 # EDIT this to how many cpu cores the node has divided by num of gpus -#SBATCH --gres=gpu:4 # EDIT this if it's not 8-gpus per node +#SBATCH --gres=gpu:8 # EDIT this if it's not 8-gpus per node #SBATCH --time=0-00:00:00 # EDIT the desired runtime #SBATCH --exclusive #SBATCH --partition=gpu # EDIT to the desired partition name From c638972ca9b74615342366f9a05915d75356ea98 Mon Sep 17 00:00:00 2001 From: srmsoumya Date: Mon, 4 Nov 2024 11:55:35 +0000 Subject: [PATCH 34/83] Modify classify, segment examples for clay v1.5 --- configs/classify_eurosat.yaml | 7 ++++--- configs/segment_chesapeake.yaml | 9 +++++---- finetune/classify/classify.py | 4 +++- finetune/classify/factory.py | 20 ++++++++++++++++---- finetune/segment/chesapeake_datamodule.py | 4 +++- finetune/segment/chesapeake_model.py | 4 ++-- finetune/segment/factory.py | 6 +++--- finetune/segment/segment.py | 6 +++++- 8 files changed, 41 insertions(+), 19 deletions(-) diff --git a/configs/classify_eurosat.yaml b/configs/classify_eurosat.yaml index 38e72eba..946a7659 100644 --- a/configs/classify_eurosat.yaml +++ b/configs/classify_eurosat.yaml @@ -2,12 +2,12 @@ seed_everything: 42 data: metadata_path: configs/metadata.yaml - batch_size: 256 + batch_size: 128 num_workers: 8 model: num_classes: 10 - ckpt_path: checkpoints/clay-v1-base.ckpt - lr: 1e-4 + ckpt_path: checkpoints/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt + lr: 5e-5 wd: 0.05 b1: 0.9 b2: 0.95 @@ -28,6 +28,7 @@ trainer: init_args: entity: developmentseed project: clay-classify + group: v1.5-test log_model: false callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint diff --git a/configs/segment_chesapeake.yaml b/configs/segment_chesapeake.yaml index d1d89dff..57e3858c 100644 --- a/configs/segment_chesapeake.yaml +++ b/configs/segment_chesapeake.yaml @@ -6,17 +6,17 @@ data: val_chip_dir: data/cvpr/ny/val/chips/ val_label_dir: data/cvpr/ny/val/labels/ metadata_path: configs/metadata.yaml - batch_size: 40 + batch_size: 16 num_workers: 8 platform: naip model: num_classes: 7 feature_maps: - - 3 - 5 - - 7 - 11 - ckpt_path: checkpoints/clay-v1-base.ckpt + - 15 + - 23 + ckpt_path: checkpoints/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-05_val-loss-0.1734.ckpt lr: 1e-5 wd: 0.05 b1: 0.9 @@ -38,6 +38,7 @@ trainer: init_args: entity: developmentseed project: clay-segment + group: v1.5-test log_model: false callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint diff --git a/finetune/classify/classify.py b/finetune/classify/classify.py index 352afd15..7b87f2e4 100644 --- a/finetune/classify/classify.py +++ b/finetune/classify/classify.py @@ -21,7 +21,9 @@ def cli_main(): """ Command-line inteface to run Clasifier model with EuroSATDataModule. """ - cli = LightningCLI(EuroSATClassifier, EuroSATDataModule) + cli = LightningCLI( + EuroSATClassifier, EuroSATDataModule, save_config_kwargs={"overwrite": True} + ) return cli diff --git a/finetune/classify/factory.py b/finetune/classify/factory.py index cebdeca2..079d3f39 100644 --- a/finetune/classify/factory.py +++ b/finetune/classify/factory.py @@ -31,20 +31,32 @@ def __init__(self, num_classes=10, ckpt_path=None): # Initialize Clay Encoder with parameters from base model. Set # mask_ratio to 0.0 & shuffle to False for downstream tasks. + # self.clay_encoder = Encoder( + # mask_ratio=0.0, + # patch_size=8, + # shuffle=False, + # dim=768, + # depth=12, + # heads=12, + # dim_head=64, + # mlp_ratio=4.0, + # ) self.clay_encoder = Encoder( mask_ratio=0.0, patch_size=8, shuffle=False, - dim=768, - depth=12, - heads=12, + dim=1024, + depth=24, + heads=16, dim_head=64, mlp_ratio=4.0, + # feature_maps=feature_maps, + # ckpt_path=ckpt_path, ) # Simple 2 layer MLP head for classification self.head = nn.Sequential( - nn.Linear(768, 512), + nn.Linear(1024, 512), nn.ReLU(), nn.Dropout(0.25), nn.Linear(512, num_classes), diff --git a/finetune/segment/chesapeake_datamodule.py b/finetune/segment/chesapeake_datamodule.py index ec7e16d0..310f2099 100644 --- a/finetune/segment/chesapeake_datamodule.py +++ b/finetune/segment/chesapeake_datamodule.py @@ -46,7 +46,9 @@ def __init__(self, chip_dir, label_dir, metadata, platform): ) # Load chip and label file names - self.chips = [chip_path.name for chip_path in self.chip_dir.glob("*.npy")] + self.chips = [chip_path.name for chip_path in self.chip_dir.glob("*.npy")][ + :1000 + ] self.labels = [re.sub("_naip-new_", "_lc_", chip) for chip in self.chips] def create_transforms(self, mean, std): diff --git a/finetune/segment/chesapeake_model.py b/finetune/segment/chesapeake_model.py index b5964ab3..949e5223 100644 --- a/finetune/segment/chesapeake_model.py +++ b/finetune/segment/chesapeake_model.py @@ -99,9 +99,9 @@ def configure_optimizers(self): ) scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, - T_0=1000, + T_0=100, T_mult=1, - eta_min=self.hparams.lr * 100, + eta_min=self.hparams.lr * 10, last_epoch=-1, ) return { diff --git a/finetune/segment/factory.py b/finetune/segment/factory.py index 0ee95db8..de439b90 100644 --- a/finetune/segment/factory.py +++ b/finetune/segment/factory.py @@ -182,9 +182,9 @@ def __init__(self, num_classes, feature_maps, ckpt_path): mask_ratio=0.0, patch_size=8, shuffle=False, - dim=768, - depth=12, - heads=12, + dim=1024, + depth=24, + heads=16, dim_head=64, mlp_ratio=4.0, feature_maps=feature_maps, diff --git a/finetune/segment/segment.py b/finetune/segment/segment.py index 7531b4d8..50b61d26 100644 --- a/finetune/segment/segment.py +++ b/finetune/segment/segment.py @@ -21,7 +21,11 @@ def cli_main(): """ Command-line inteface to run Segmentation Model with ChesapeakeDataModule. """ - cli = LightningCLI(ChesapeakeSegmentor, ChesapeakeDataModule) + cli = LightningCLI( + ChesapeakeSegmentor, + ChesapeakeDataModule, + save_config_kwargs={"overwrite": True}, + ) return cli From 05c167e3d8c55c8b37e4fc7637b5883c2e280bc2 Mon Sep 17 00:00:00 2001 From: Soumya Ranjan Mohanty Date: Mon, 4 Nov 2024 17:30:07 +0530 Subject: [PATCH 35/83] Check non MRL loss for the model (#331) - Use DINO v2 with linear projection - Don't load from base checkpoint of MRL --- configs/config.yaml | 12 ++++++------ src/callbacks_wandb.py | 12 +++--------- src/datamodule.py | 2 +- src/model.py | 23 +++++++++++++++++------ src/module.py | 21 +++++++++++++++++++++ src/mrl.py | 5 +++-- train_clay_v2.sh | 8 ++++++-- train_environment.yml | 8 +++++--- trainer.py | 4 +--- 9 files changed, 63 insertions(+), 32 deletions(-) diff --git a/configs/config.yaml b/configs/config.yaml index 1f9de9ee..aad3353f 100644 --- a/configs/config.yaml +++ b/configs/config.yaml @@ -20,7 +20,7 @@ model: patch_size: 8 shuffle: True metadata_path: configs/metadata.yaml - teacher: samvit_base_patch16.sa1b + teacher: vit_large_patch14_reg4_dinov2.lvd142m dolls: [16, 32, 64, 128, 256, 768, 1024] doll_weights: [1, 1, 1, 1, 1, 1, 1] lr: 5e-6 @@ -32,7 +32,7 @@ trainer: accelerator: gpu strategy: ddp devices: 8 - num_nodes: 20 + num_nodes: 48 precision: bf16-mixed log_every_n_steps: 1 max_epochs: 1000 @@ -48,9 +48,9 @@ trainer: init_args: entity: developmentseed project: clay - group: v1.5 - # id: v8jh2pn9 - # resume: must + group: v1.5-nomrl-dinov2 + id: 0uy3in7l + resume: must log_model: false callbacks: - class_path: lightning.pytorch.callbacks.ModelCheckpoint @@ -70,4 +70,4 @@ trainer: - class_path: src.callbacks_wandb.LogIntermediatePredictions plugins: - class_path: lightning.pytorch.plugins.io.AsyncCheckpointIO -ckpt_path: null +ckpt_path: checkpoints/v1.5.0/last.ckpt diff --git a/src/callbacks_wandb.py b/src/callbacks_wandb.py index 15ec437c..374867fc 100644 --- a/src/callbacks_wandb.py +++ b/src/callbacks_wandb.py @@ -257,9 +257,7 @@ def on_validation_end( for j in range(n_cols): # Plot actual images in rows 0 and 2 - axs[0, j].imshow( - batch["pixels"][j][0], cmap="viridis" - ) + axs[0, j].imshow(batch["pixels"][j][0], cmap="viridis") axs[0, j].set_title(f"Actual {j}") axs[0, j].axis("off") @@ -271,15 +269,11 @@ def on_validation_end( axs[2, j].axis("off") # Plot predicted images in rows 1 and 3 - axs[1, j].imshow( - pixels[j][0], cmap="viridis" - ) + axs[1, j].imshow(pixels[j][0], cmap="viridis") axs[1, j].set_title(f"Pred {j}") axs[1, j].axis("off") - axs[3, j].imshow( - pixels[j + n_cols][0], cmap="viridis" - ) + axs[3, j].imshow(pixels[j + n_cols][0], cmap="viridis") axs[3, j].set_title(f"Pred {j+n_cols}") axs[3, j].axis("off") diff --git a/src/datamodule.py b/src/datamodule.py index ac7d7b5a..dc6c0901 100644 --- a/src/datamodule.py +++ b/src/datamodule.py @@ -253,7 +253,7 @@ def setup(self, stage: Literal["fit", "predict"] | None = None) -> None: # chips_path = list(dp.list_files_by_s3(masks="*.npz")) # else: # if self.data_dir is a local data path chips_path = sorted(list(Path(self.data_dir).glob("**/*.npz"))) - chips_platform = [chip.parent.parent.name for chip in chips_path] + chips_platform = [chip.parent.name for chip in chips_path] # chips_platform = [chip.parent.parent.name for chip in chips_path] print(f"Total number of chips: {len(chips_path)}") diff --git a/src/model.py b/src/model.py index 797590d7..9648002f 100644 --- a/src/model.py +++ b/src/model.py @@ -7,10 +7,10 @@ import torch.nn.functional as F from einops import rearrange, reduce, repeat from torch import nn +from torchvision.transforms import v2 from src.backbone import Transformer from src.factory import DynamicEmbedding -from src.mrl import MRL, MRLLoss from src.utils import posemb_sincos_2d_with_gsd torch.set_float32_matmul_precision("medium") @@ -386,8 +386,13 @@ def __init__( # noqa: PLR0913 self.shuffle = shuffle self.metadata = metadata self.teacher = timm.create_model(teacher, pretrained=True, num_classes=0) - self.mrl = MRL(features=self.teacher.num_features, dolls=dolls) - self.mrl_loss = MRLLoss(weights=doll_weights) + self.teacher_chip_size = 518 + self.teacher_resize = v2.Resize( + size=(self.teacher_chip_size, self.teacher_chip_size) + ) + # self.mrl = MRL(features=self.teacher.num_features, dolls=dolls) + # self.mrl_loss = MRLLoss(weights=doll_weights) + self.proj = nn.Linear(dim, self.teacher.num_features) self.encoder = Encoder( mask_ratio=mask_ratio, @@ -516,8 +521,11 @@ def forward(self, datacube): if platform == "modis": reconstruction_loss /= 10 - # MRL - representations = self.mrl(encoded_unmasked_patches[:, 0, :]) # [(B D') ...] + # # MRL + # representations = self.mrl(encoded_unmasked_patches[:, 0, :]) # [(B D') ...] + + # PROJ + representations = self.proj(encoded_unmasked_patches[:, 0, :]) # [B D'] with torch.no_grad(): if platform == "sentinel-1-rtc": @@ -529,9 +537,12 @@ def forward(self, datacube): # Read RGB bands from the sensor to feed the teacher model indices = self.metadata[platform].rgb_indices rgb = datacube["pixels"][:, indices, :, :] + rgb = self.teacher_resize(rgb) target = self.teacher(rgb) + # target = self.teacher(rgb) - representation_loss = self.mrl_loss(representations, target) + # representation_loss = self.mrl_loss(representations, target) + representation_loss = 1.0 - F.cosine_similarity(representations, target).mean() loss = 0.9 * reconstruction_loss + 0.1 * representation_loss return (loss, reconstruction_loss, representation_loss) diff --git a/src/module.py b/src/module.py index 39b9df3a..eafcf815 100644 --- a/src/module.py +++ b/src/module.py @@ -27,6 +27,7 @@ def __init__( # noqa: PLR0913 embeddings_level: Literal["mean", "patch", "group"] = "mean", ): super().__init__() + # self.strict_loading = False # Allow partial loading to check if MRL was the bug self.save_hyperparameters(logger=True) self.metadata = Box(yaml.safe_load(open(metadata_path))) model_map = { @@ -47,6 +48,26 @@ def __init__( # noqa: PLR0913 "doll_weights": doll_weights, } self.model = model_map[model_size](**model_args) + # checkpoint_path = 'mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt' + # checkpoint = torch.load(checkpoint_path, map_location="cpu") + # # Extract the state dictionary + # state_dict = checkpoint['state_dict'] + + # # Modify the state dictionary + # new_state_dict = OrderedDict() + # for k, v in state_dict.items(): + # # Remove 'model.' prefix if it exists + # if k.startswith('model.'): + # k = k[len('model.'):] + # # Exclude keys related to the 'teacher' + # if not (k.startswith('teacher') or k.startswith('mrl')): + # new_state_dict[k] = v + # with torch.no_grad(): + # # Load the modified state dictionary into your model + # missing_keys, unexpected_keys = self.model.load_state_dict(new_state_dict, strict=False) + # # Optionally, print missing and unexpected keys + # print(f"Missing keys: {missing_keys}") + # print(f"Unexpected keys: {unexpected_keys}") else: raise ValueError( f"Invalid model size {model_size}. Expected one of {model_map.keys()}" diff --git a/src/mrl.py b/src/mrl.py index a7202ad4..12ee22e7 100644 --- a/src/mrl.py +++ b/src/mrl.py @@ -9,12 +9,13 @@ class MRL(nn.Module): def __init__(self, features, dolls: list = [16, 32, 64, 128, 256, 768]) -> None: super().__init__() self.dolls = dolls + self.layers = nn.ModuleDict() for doll in dolls: - setattr(self, f"mrl_{doll}", nn.Linear(doll, features)) + self.layers[f"mrl_{doll}"] = nn.Linear(doll, features) def forward(self, x): "x: (batch, features)" - logits = [getattr(self, f"mrl_{doll}")(x[:, :doll]) for doll in self.dolls] + logits = [self.layers[f"mrl_{doll}"](x[:, :doll]) for doll in self.dolls] return logits diff --git a/train_clay_v2.sh b/train_clay_v2.sh index b39332e3..f680d7bc 100644 --- a/train_clay_v2.sh +++ b/train_clay_v2.sh @@ -1,13 +1,14 @@ #!/bin/bash #SBATCH --job-name=clay-laucher -#SBATCH --nodes=20 +#SBATCH --nodes=24 #SBATCH --ntasks-per-node=8 # EDIT if it's not 8-gpus per node #SBATCH --cpus-per-task=12 # EDIT this to how many cpu cores the node has divided by num of gpus #SBATCH --gres=gpu:8 # EDIT this if it's not 8-gpus per node #SBATCH --time=0-00:00:00 # EDIT the desired runtime #SBATCH --exclusive #SBATCH --partition=gpu # EDIT to the desired partition name +#SBATCH --nodelist=gpu-dy-g6-[1-12],gpu-dy-g5-[1-12] #SBATCH --output=%x-%j-%N.out echo "START TIME: $(date)" @@ -31,8 +32,11 @@ LOG_PATH="main_log.txt" # PTL doesn't need a special launcher LAUNCHER="python -u" +# Capture the number of nodes allocated by Slurm +NUM_NODES=$SLURM_JOB_NUM_NODES + # EDIT the path+name of the python script and whatever args it needs -PROGRAM="trainer.py fit --config configs/config.yaml" +PROGRAM="trainer.py fit --config configs/config.yaml --trainer.num_nodes=$NUM_NODES" export CMD="$LAUNCHER $PROGRAM" diff --git a/train_environment.yml b/train_environment.yml index 433e5b6c..78923977 100644 --- a/train_environment.yml +++ b/train_environment.yml @@ -1,21 +1,23 @@ name: claymodel channels: - - pytorch - conda-forge + - nvidia + - pytorch dependencies: - python=3.11 - pip - pip: + - --extra-index-url https://download.pytorch.org/whl/cu121 + - torch==2.4.0+cu121 + - torchvision==0.19.0+cu121 - einops~=0.7.0 - geopandas - jsonargparse[signatures]>=4.27.7 - lightning - matplotlib - python-box - - torch - scikit-image - scikit-learn - timm - - torchvision - vit-pytorch - wandb diff --git a/trainer.py b/trainer.py index cce9ecac..509925fc 100644 --- a/trainer.py +++ b/trainer.py @@ -22,9 +22,7 @@ def cli_main(): Command-line inteface to run ClayMAE with ClayDataModule. """ cli = LightningCLI( - ClayMAEModule, - ClayDataModule, - save_config_kwargs={"overwrite": True} + ClayMAEModule, ClayDataModule, save_config_kwargs={"overwrite": True} ) return cli From 38a3c27e1a28e1d0a83b114cc1d8c7e6c8d9eaab Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Sun, 28 Jul 2024 23:08:38 +0100 Subject: [PATCH 36/83] Document MODIS data sampling --- docs/release-notes/data_sampling.md | 62 +++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/docs/release-notes/data_sampling.md b/docs/release-notes/data_sampling.md index 4741064a..c3650964 100644 --- a/docs/release-notes/data_sampling.md +++ b/docs/release-notes/data_sampling.md @@ -112,6 +112,67 @@ and a maximum of 2000 scenes for each catalog that was included. We selected the latest imagery for each of the available regions of new zealand. The list of catalogs is in the linz processor file. +### MODIS sampling strategy + +For MODIS we used the [Surface Reflectance 8-Day (500m)](https://planetarycomputer.microsoft.com/dataset/modis-09A1-061) +product. The data is distributed in SIN grid tiles. We included all SIN grid +tiles that do not have any nodata inside. The selected SIN grid tiles are then +transform to EPSG:3857 for all tiles. This results in some variation between the +nominal resolution, although the original resolution from the SIN projection is +500 meters. For input to the model, we assumed the 500m resolution as a fixed +resolution size for all tiles. + +Algorithm to determine which tiles do not have nodata is shown in the code block +below. This resulted in 233 SIN grid tiles to be selected. For each of these +we sampled the first STAC search result for each month in each year from 2018 +until 2023. This therefore resulted in 72 (`6 years * 12 months`) separate scenes +for each of the 233 SIN grid tiles. + +Script for selection of SIN grid tiles included in the sampling: + +```python +from multiprocessing import Pool +import rasterio +import planetary_computer as pc +import pystac_client +import numpy as np + +SIN_GRID_TILES = [] +for i in SIN_VERTICAL_RANGE: + for j in SIN_HORIZONTAL_RANGE: + SIN_GRID_TILES.append((i, j)) + +def evaluate_nodata(i, j): + catalog = pystac_client.Client.open(STAC_API, modifier=pc.sign_inplace) + items = catalog.search( + collections=[COLLECTION], + query={ + "modis:vertical-tile": { + "eq": i, + }, + "modis:horizontal-tile": { + "eq": j, + }, + }, + max_items=1, + ) + item = list(items.item_collection())[0] + + with rasterio.open(item.assets["sur_refl_b01"].href) as src: + data = src.read() + + nodata = np.sum(data == -28672) + + if nodata == 0: + print(i, j) + return i, j + +if __name__ == '__main__': + with Pool(16) as p: + indexes = p.starmap(evaluate_nodata, SIN_GRID_TILES) + print("done") + print(indexes) +``` ## Data preparation @@ -136,6 +197,7 @@ Using stacchip, we created a dataset with a size of 33.8 TB of imagery, with abo | Landsat-c2l1 | 5827333 | | Landsat-c2l2-sr | 5790651 | | Sentinel-1-rtc | 16133394 | +| MODIS | 1350864 | # Older versions From 4f464b926bb4a0e2c7203c34c3360b9011417ebf Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Thu, 26 Sep 2024 08:39:16 +0100 Subject: [PATCH 37/83] intermediate commit --- docker/Dockerfile | 11 +++ docker/all-naip.py | 236 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 247 insertions(+) create mode 100644 docker/Dockerfile create mode 100644 docker/all-naip.py diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 00000000..add25a89 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,11 @@ +FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 + +RUN pip install stacchip + +WORKDIR /src + +ADD data data +ADD checkpoints checkpoints +ADD all-naip.py . + +ENTRYPOINT [ "python", "all-naip.py" ] diff --git a/docker/all-naip.py b/docker/all-naip.py new file mode 100644 index 00000000..6fb04272 --- /dev/null +++ b/docker/all-naip.py @@ -0,0 +1,236 @@ +import datetime +import io +import math +import os +import tempfile +import zipfile + +import boto3 +import geoarrow.pyarrow as ga +import numpy as np +import pyarrow as pa +import torch +from geoarrow.pyarrow import io as gaio +from rasterio.errors import RasterioIOError +from rio_stac import create_stac_item +from stacchip.chipper import Chipper +from stacchip.indexer import NoStatsChipIndexer +from torchvision.transforms import v2 + + +def normalize_timestamp(date): + week = date.isocalendar().week * 2 * np.pi / 52 + hour = date.hour * 2 * np.pi / 24 + + return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour)) + + +def normalize_latlon(lat, lon): + lat = lat * np.pi / 180 + lon = lon * np.pi / 180 + + return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon)) + + +def prepare_datacube(datetimes, bboxs, pixels, gsd): + # Set mean, std, and wavelengths metadata + mean = [ + 110.16, + 115.41, + 98.15, + 139.04, + ] + std = [47.23, 39.82, 35.43, 49.86] + waves = [0.65, 0.56, 0.48, 0.842] + + transform = v2.Compose( + [ + v2.Normalize(mean=mean, std=std), + ] + ) + + times = [normalize_timestamp(dat) for dat in datetimes] + week_norm = [dat[0] for dat in times] + hour_norm = [dat[1] for dat in times] + time_norm = np.hstack((week_norm, hour_norm)) + + latlons = [normalize_latlon(*bbox.centroid.coords[0]) for bbox in bboxs] + lat_norm = [dat[0] for dat in latlons] + lon_norm = [dat[1] for dat in latlons] + latlon_norm = np.hstack((lat_norm, lon_norm)) + + gsd = [gsd] + + pixels_norm = transform(pixels) + + return waves, time_norm, latlon_norm, gsd, pixels_norm + + +def get_pixels(item): + indexer = NoStatsChipIndexer(item) + + # Instanciate the chipper + chipper = Chipper(indexer) + + # Get first chip for the "image" asset key + chips = [] + datetimes = [] + bboxs = [] + chip_ids = [] + item_ids = [] + for idx, (x, y, chip) in enumerate(chipper): + chips.append(chip) + datetimes.append(item.datetime) + bboxs.append(indexer.get_chip_bbox(x, y)) + chip_ids.append((x, y)) + item_ids.append(item.id) + + pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips]) + return bboxs, datetimes, pixels + + +def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchsize): # noqa: PLR0913 + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + print(f"Using device {device}") + # Run the clay encoder + embeddings = None + for i in range(0, len(pixels_norm), batchsize): + if i % 500 == 0: + print(f"Iteration {i}") + datacube = { + "pixels": torch.tensor( + pixels_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "time": torch.tensor( + time_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "latlon": torch.tensor( + latlon_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "waves": torch.tensor(waves, dtype=torch.float32, device=device), + "gsd": torch.tensor(gsd, dtype=torch.float32, device=device), + } + with torch.no_grad(): + cls_embedding = clay(datacube) + if embeddings is None: + embeddings = cls_embedding + else: + embeddings = torch.vstack((embeddings, cls_embedding)) + + return embeddings + + +def open_scene_list(): + with zipfile.ZipFile("data/naip-manifest.txt.zip") as zf: + with io.TextIOWrapper(zf.open("naip-manifest.txt"), encoding="utf-8") as f: + data = f.readlines() + data = [dat.rstrip() for dat in data if "rgbir_cog" in dat] + return data + + +def load_clay(): + if torch.cuda.is_available(): + checkpoint = "checkpoints/clay-v1-encoder.pt2" + else: + checkpoint = "checkpoints/clay-v1-encoder-cpu.pt2" + + return torch.export.load(checkpoint).module() + + +def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path, item_id): # noqa: PLR0913 + index = { + "embeddings": [np.ascontiguousarray(dat) for dat in embeddings.cpu().numpy()], + "geometry": ga.as_geoarrow([dat.wkt for dat in bboxs]), + } + + table = pa.table( + index, + metadata={ + "date": datestr, + "gsd": str(gsd[0]), + "uri": f"s3://naip-analytic/{path}", + }, + ) + + writer = pa.BufferOutputStream() + gaio.write_geoparquet_table(table, writer) + body = bytes(writer.getvalue()) + s3_resource = boto3.resource("s3") + s3_bucket = s3_resource.Bucket(name=destination_bucket) + s3_bucket.put_object( + Body=body, + Key=f"{item_id}.parquet", + ) + + +def process_scene(clay, path, destination_bucket, batchsize): + state = path.split("/")[0] + datestr = path.split("/")[-1].split("_")[-1].split(".txt")[0] + gsd = float(path.split("/")[2].replace("cm", "")) / 100 + date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) + print(f"Processing {path} in state {state} and date {date}") + + with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as f: + s3 = boto3.client("s3") + s3.download_fileobj( + "naip-analytic", path, f, ExtraArgs={"RequestPayer": "requester"} + ) + + item = create_stac_item(f.name, with_proj=True) + item.datetime = date + item.id = f"{state}_{path.split('/')[-1].replace('.tif', '')}" + + try: + bboxs, datetimes, pixels = get_pixels(item) + except RasterioIOError: + print("Skipping scene due to rasterio io error") + return + + waves, time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( + datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd + ) + + embeddings = get_embeddings( + clay=clay, + pixels_norm=pixels_norm, + time_norm=time_norm, + latlon_norm=latlon_norm, + waves=waves, + gsd=gsd, + batchsize=batchsize, + ) + + write_to_table( + embeddings=embeddings, + bboxs=bboxs, + datestr=datestr, + gsd=gsd, + destination_bucket=destination_bucket, + path=path, + item_id=item.id, + ) + + +def process(): + if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: + raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") + index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) + items_per_job = int(os.environ.get("ITEMS_PER_JOB", 100)) + batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) + destination_bucket = "clay-v1-naip-embeddings" + + scenes = open_scene_list() + clay = load_clay() + + for i in range(index * items_per_job, (index + 1) * items_per_job): + process_scene( + clay=clay, + path=scenes[i], + destination_bucket=destination_bucket, + batchsize=batchsize, + ) + + +if __name__ == "__main__": + process() + print("Done!") From 8ecff64b54e9b6af28c17b4cd450de783463747a Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 30 Sep 2024 14:55:26 +0100 Subject: [PATCH 38/83] intermediate --- docker/Dockerfile | 9 ++++++--- docker/all-naip.py | 16 +++++++++++----- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index add25a89..ff473c1e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -4,8 +4,11 @@ RUN pip install stacchip WORKDIR /src -ADD data data -ADD checkpoints checkpoints -ADD all-naip.py . +ADD docker/data data +ADD docker/checkpoints checkpoints +ADD docker/all-naip.py . + +ADD environment.yml environment.yml +RUN conda install -y environment.yml ENTRYPOINT [ "python", "all-naip.py" ] diff --git a/docker/all-naip.py b/docker/all-naip.py index 6fb04272..63fadb85 100644 --- a/docker/all-naip.py +++ b/docker/all-naip.py @@ -1,5 +1,6 @@ import datetime import io +import logging import math import os import tempfile @@ -17,6 +18,10 @@ from stacchip.indexer import NoStatsChipIndexer from torchvision.transforms import v2 +logging.basicConfig() +logger = logging.getLogger("clay") +logger.setLevel(logging.DEBUG) + def normalize_timestamp(date): week = date.isocalendar().week * 2 * np.pi / 52 @@ -91,12 +96,12 @@ def get_pixels(item): def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchsize): # noqa: PLR0913 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - print(f"Using device {device}") + logger.debug(f"Using device {device}") # Run the clay encoder embeddings = None for i in range(0, len(pixels_norm), batchsize): if i % 500 == 0: - print(f"Iteration {i}") + logger.debug(f"Iteration {i}") datacube = { "pixels": torch.tensor( pixels_norm[i : (i + batchsize)], dtype=torch.float32, device=device @@ -168,7 +173,7 @@ def process_scene(clay, path, destination_bucket, batchsize): datestr = path.split("/")[-1].split("_")[-1].split(".txt")[0] gsd = float(path.split("/")[2].replace("cm", "")) / 100 date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) - print(f"Processing {path} in state {state} and date {date}") + logger.debug(f"Processing {path} in state {state} and date {date}") with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as f: s3 = boto3.client("s3") @@ -183,7 +188,7 @@ def process_scene(clay, path, destination_bucket, batchsize): try: bboxs, datetimes, pixels = get_pixels(item) except RasterioIOError: - print("Skipping scene due to rasterio io error") + logger.debug("Skipping scene due to rasterio io error") return waves, time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( @@ -232,5 +237,6 @@ def process(): if __name__ == "__main__": + logger.debug("Starting") process() - print("Done!") + logger.debug("Done!") From e701052d1acb468c501cb71d881813745481e572 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 1 Oct 2024 10:42:25 +0100 Subject: [PATCH 39/83] Update to v1.5 and add logging --- docker/Dockerfile | 18 ++++++++++-------- docker/all-naip.py | 25 ++++++++++++++++++------- 2 files changed, 28 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index ff473c1e..9fa8d4a2 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,14 +1,16 @@ FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 -RUN pip install stacchip +WORKDIR /code -WORKDIR /src +ADD src src -ADD docker/data data -ADD docker/checkpoints checkpoints -ADD docker/all-naip.py . +ADD docker/environment.yml docker-environment.yml +RUN mamba env create --file docker-environment.yml + +ADD data/naip-manifest.txt.zip /data/naip-manifest.txt.zip +ADD data/checkpoints/clay-model-v1.5.0-september-30.ckpt /data/clay-model-v1.5.0-september-30.ckpt -ADD environment.yml environment.yml -RUN conda install -y environment.yml +ADD docker/all-naip.py . +ADD configs configs -ENTRYPOINT [ "python", "all-naip.py" ] +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] diff --git a/docker/all-naip.py b/docker/all-naip.py index 63fadb85..43a021bb 100644 --- a/docker/all-naip.py +++ b/docker/all-naip.py @@ -18,6 +18,8 @@ from stacchip.indexer import NoStatsChipIndexer from torchvision.transforms import v2 +from src.module import ClayMAEModule + logging.basicConfig() logger = logging.getLogger("clay") logger.setLevel(logging.DEBUG) @@ -96,7 +98,7 @@ def get_pixels(item): def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchsize): # noqa: PLR0913 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - logger.debug(f"Using device {device}") + logger.debug(f"Using device {device} to create {len(pixels_norm)} embeddings") # Run the clay encoder embeddings = None for i in range(0, len(pixels_norm), batchsize): @@ -114,6 +116,7 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs ), "waves": torch.tensor(waves, dtype=torch.float32, device=device), "gsd": torch.tensor(gsd, dtype=torch.float32, device=device), + "platform": ["naip"], } with torch.no_grad(): cls_embedding = clay(datacube) @@ -126,20 +129,28 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs def open_scene_list(): - with zipfile.ZipFile("data/naip-manifest.txt.zip") as zf: + with zipfile.ZipFile("/data/naip-manifest.txt.zip") as zf: with io.TextIOWrapper(zf.open("naip-manifest.txt"), encoding="utf-8") as f: data = f.readlines() data = [dat.rstrip() for dat in data if "rgbir_cog" in dat] + logger.debug(f"Found {len(data)} NAIP scenes in manifest") return data def load_clay(): - if torch.cuda.is_available(): - checkpoint = "checkpoints/clay-v1-encoder.pt2" - else: - checkpoint = "checkpoints/clay-v1-encoder-cpu.pt2" + device = "cuda" if torch.cuda.is_available() else "cpu" + model = ClayMAEModule.load_from_checkpoint( + checkpoint_path="/data/clay-model-v1.5.0-september-30.ckpt", + metadata_path="configs/metadata.yaml", + model_size="large", + dolls=[16, 32, 64, 128, 256, 768, 1024], + doll_weights=[1, 1, 1, 1, 1, 1, 1], + mask_ratio=0.0, + shuffle=False, + ) + model.eval() - return torch.export.load(checkpoint).module() + return model.to(device) def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path, item_id): # noqa: PLR0913 From d46522a349567e0fd8ce17955b369fb68bdcbcd0 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 1 Oct 2024 14:14:56 +0100 Subject: [PATCH 40/83] Fix for v1.5 module input --- docker/all-naip.py | 44 ++++++++++++++++++++++++++------------------ 1 file changed, 26 insertions(+), 18 deletions(-) diff --git a/docker/all-naip.py b/docker/all-naip.py index 43a021bb..bd1eddbe 100644 --- a/docker/all-naip.py +++ b/docker/all-naip.py @@ -5,6 +5,7 @@ import os import tempfile import zipfile +from pathlib import Path import boto3 import geoarrow.pyarrow as ga @@ -20,6 +21,9 @@ from src.module import ClayMAEModule +MANIFEST = "/data/naip-manifest.txt.zip" +CHECKPOINT = "/data/clay-model-v1.5.0-september-30.ckpt" + logging.basicConfig() logger = logging.getLogger("clay") logger.setLevel(logging.DEBUG) @@ -102,8 +106,9 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs # Run the clay encoder embeddings = None for i in range(0, len(pixels_norm), batchsize): - if i % 500 == 0: + if i / batchsize % 5 == 0: logger.debug(f"Iteration {i}") + datacube = { "pixels": torch.tensor( pixels_norm[i : (i + batchsize)], dtype=torch.float32, device=device @@ -119,20 +124,24 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs "platform": ["naip"], } with torch.no_grad(): - cls_embedding = clay(datacube) + unmsk_patch, unmsk_idx, msk_idx, msk_matrix = clay.model.encoder(datacube) + # The first embedding is the class token, which is the + # overall single embedding we want to keep. + cls_embeddings = unmsk_patch[:, 0, :] if embeddings is None: - embeddings = cls_embedding + embeddings = cls_embeddings else: - embeddings = torch.vstack((embeddings, cls_embedding)) + embeddings = torch.vstack((embeddings, cls_embeddings)) return embeddings def open_scene_list(): - with zipfile.ZipFile("/data/naip-manifest.txt.zip") as zf: + with zipfile.ZipFile(MANIFEST) as zf: with io.TextIOWrapper(zf.open("naip-manifest.txt"), encoding="utf-8") as f: data = f.readlines() - data = [dat.rstrip() for dat in data if "rgbir_cog" in dat] + data = [Path(dat.rstrip()) for dat in data if "rgbir_cog"] + data = [dat for dat in data if dat.suffix == ".tif"] logger.debug(f"Found {len(data)} NAIP scenes in manifest") return data @@ -140,7 +149,7 @@ def open_scene_list(): def load_clay(): device = "cuda" if torch.cuda.is_available() else "cpu" model = ClayMAEModule.load_from_checkpoint( - checkpoint_path="/data/clay-model-v1.5.0-september-30.ckpt", + checkpoint_path=CHECKPOINT, metadata_path="configs/metadata.yaml", model_size="large", dolls=[16, 32, 64, 128, 256, 768, 1024], @@ -153,7 +162,7 @@ def load_clay(): return model.to(device) -def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path, item_id): # noqa: PLR0913 +def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): index = { "embeddings": [np.ascontiguousarray(dat) for dat in embeddings.cpu().numpy()], "geometry": ga.as_geoarrow([dat.wkt for dat in bboxs]), @@ -175,31 +184,31 @@ def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path, it s3_bucket = s3_resource.Bucket(name=destination_bucket) s3_bucket.put_object( Body=body, - Key=f"{item_id}.parquet", + Key=f"{path.parent}/{path.stem}.parquet", ) def process_scene(clay, path, destination_bucket, batchsize): - state = path.split("/")[0] - datestr = path.split("/")[-1].split("_")[-1].split(".txt")[0] - gsd = float(path.split("/")[2].replace("cm", "")) / 100 + state = path.parts[0] + datestr = path.stem.split("_")[-1] date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) + gsd = float(path.parts[2].replace("cm", "")) / 100 logger.debug(f"Processing {path} in state {state} and date {date}") - with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as f: + with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as fl: s3 = boto3.client("s3") s3.download_fileobj( - "naip-analytic", path, f, ExtraArgs={"RequestPayer": "requester"} + "naip-analytic", str(path), fl, ExtraArgs={"RequestPayer": "requester"} ) - item = create_stac_item(f.name, with_proj=True) + item = create_stac_item(fl.name, with_proj=True) item.datetime = date - item.id = f"{state}_{path.split('/')[-1].replace('.tif', '')}" + item.id = f"{state}_{path.stem}" try: bboxs, datetimes, pixels = get_pixels(item) except RasterioIOError: - logger.debug("Skipping scene due to rasterio io error") + logger.warning("Skipping scene due to rasterio io error") return waves, time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( @@ -223,7 +232,6 @@ def process_scene(clay, path, destination_bucket, batchsize): gsd=gsd, destination_bucket=destination_bucket, path=path, - item_id=item.id, ) From c7283d77c7c779d3e66451e0f7a62ce053775741 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 22:55:31 +0100 Subject: [PATCH 41/83] Update adding sentinel-2 --- embeddings/Dockerfile | 18 ++++ embeddings/all-naip.py | 131 +++++++++++++++++++++++++++++ embeddings/all-sentinel.py | 104 +++++++++++++++++++++++ embeddings/environment.yml | 40 +++++++++ embeddings/strategy.sh | 47 +++++++++++ embeddings/utils.py | 168 +++++++++++++++++++++++++++++++++++++ 6 files changed, 508 insertions(+) create mode 100644 embeddings/Dockerfile create mode 100644 embeddings/all-naip.py create mode 100644 embeddings/all-sentinel.py create mode 100644 embeddings/environment.yml create mode 100644 embeddings/strategy.sh create mode 100644 embeddings/utils.py diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile new file mode 100644 index 00000000..3e32308e --- /dev/null +++ b/embeddings/Dockerfile @@ -0,0 +1,18 @@ +FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 + +WORKDIR /model + +RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt /data/clay-model-v1.5.0-october-12.ckpt + +RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model + +# ADD docker/environment.yml docker-environment.yml +RUN mamba env create --file environment.yml + +ADD data/naip-manifest.txt.zip /data/naip-manifest.txt.zip +ADD data/checkpoints/clay-model-v1.5.0-september-30.ckpt /data/clay-model-v1.5.0-september-30.ckpt + +ADD docker/all-naip.py . +ADD configs configs + +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py new file mode 100644 index 00000000..044a6b5c --- /dev/null +++ b/embeddings/all-naip.py @@ -0,0 +1,131 @@ +import datetime +import io +import logging +import os +import tempfile +import zipfile +from pathlib import Path + +import boto3 +from rasterio.errors import RasterioIOError +from rio_stac import create_stac_item +from stacchip.indexer import NoStatsChipIndexer + +from embeddings.utils import ( + get_embeddings, + get_pixels, + load_clay, + prepare_datacube, + write_to_table, +) + +logging.basicConfig() +logger = logging.getLogger("clay") +logger.setLevel(logging.DEBUG) + + +# MANIFEST = "/data/naip-manifest.txt.zip" +MANIFEST = "/Users/tam/Documents/repos/model/data/naip-manifest.txt.zip" +MEAN = [ + 110.16, + 115.41, + 98.15, + 139.04, +] +STD = [47.23, 39.82, 35.43, 49.86] +WAVES = [0.65, 0.56, 0.48, 0.842] +EMBEDDINGS_BUCKET = "clay-embeddings-naip" + + +def open_scene_list(): + """ + Read the naip-analytic manifest file and extract a list of NAIP + scenes as tif files to process. + + The file used here is the zipped version of the original manifest file. + """ + with zipfile.ZipFile(MANIFEST) as zf: + with io.TextIOWrapper(zf.open("naip-manifest.txt"), encoding="utf-8") as f: + data = f.readlines() + data = [Path(dat.rstrip()) for dat in data if "rgbir_cog"] + data = [dat for dat in data if dat.suffix == ".tif"] + logger.debug(f"Found {len(data)} NAIP scenes in manifest") + return data + + +def process_scene(clay, path, batchsize): + """ + Embeds a slingle NAIP scene. + """ + state = path.parts[0] + datestr = path.stem.split("_")[-1] + date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) + gsd = float(path.parts[2].replace("cm", "")) / 100 + logger.debug(f"Processing {path} in state {state} and date {date}") + + with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as fl: + s3 = boto3.client("s3") + s3.download_fileobj( + "naip-analytic", str(path), fl, ExtraArgs={"RequestPayer": "requester"} + ) + + item = create_stac_item(fl.name, with_proj=True) + + item.datetime = date + item.id = f"{state}_{path.stem}" + + try: + bboxs, datetimes, pixels = get_pixels(item, NoStatsChipIndexer) + except RasterioIOError: + logger.warning("Skipping scene due to rasterio io error") + return + + time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( + mean=MEAN, std=STD, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd + ) + # Embed data + cls_embeddings, patch_embeddings = get_embeddings( + clay=clay, + pixels_norm=pixels_norm, + time_norm=time_norm, + latlon_norm=latlon_norm, + waves=WAVES, + gsd=gsd, + batchsize=batchsize, + ) + # Write class embeddings + + kwargs = dict( + bboxs=bboxs, + datestr=datestr, + gsd=gsd, + destination_bucket=EMBEDDINGS_BUCKET, + path=path, + ) + write_to_table(embeddings=cls_embeddings, **kwargs) + # Write patch embeddings + write_to_table(embeddings=patch_embeddings, **kwargs) + + +def process(): + if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: + raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") + index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) + items_per_job = int(os.environ.get("ITEMS_PER_JOB", 2)) + batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 2)) + + scenes = open_scene_list() + clay = load_clay() + + for i in range(index * items_per_job, (index + 1) * items_per_job): + process_scene( + clay=clay, + path=scenes[i], + batchsize=batchsize, + ) + + +if __name__ == "__main__": + logger.debug("Starting") + process() + logger.debug("Done!") diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py new file mode 100644 index 00000000..e9c6fc1a --- /dev/null +++ b/embeddings/all-sentinel.py @@ -0,0 +1,104 @@ +import gzip +import json +import logging +import os +import sys + +import boto3 +from pystac import Item +from stacchip.indexer import Sentinel2Indexer + +from embeddings.utils import load_clay + +logger = logging.getLogger("clay") + + +# Create 2023 file from full archive + +s3 = boto3.resource("s3") + + +def log(comment): + sys.stdout.write(f"\r{comment}") + + +log_every = 100000 + +count = 0 + +# with gzip.open(f"data/element84-tiles-2023.gz", "wt") as dst: +# with gzip.open(f"data/element84-tiles.list.gz") as fl: +# line = fl.readline() +# while line: +# line = line.decode().rstrip() +# c = line.split("/") +# if c[4] == "2019": +# line = fl.readline() +# continue +# elif int(c[7]) < 2023: +# line = fl.readline() +# continue +# elif not line.endswith("L2A.json"): +# line = fl.readline() +# continue + +# count += 1 +# if count % log_every == 0: +# log(f"Found {count} scenes... {line}") + +# dst.write(line + "\n") +# line = fl.readline() + + +def open_manifest(path): + pass + + +data = open_manifest() + +index = 42 + +key = data[index].decode().rstrip().replace("s3://sentinel-cogs/", "") + +stac_json = json.load(s3.Object("sentinel-cogs", key).get()["Body"]) + +item = Item.from_dict(stac_json) + + +indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) + + +SCENES_LIST = "data/element84-tiles-2023.gz" + + +def open_scenes_list(): + with gzip.open(SCENES_LIST) as fl: + return fl.readlines() + + +def process_scene(clay, path, batchsize): + pass + + +def process(): + if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: + raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") + index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) + items_per_job = int(os.environ.get("ITEMS_PER_JOB", 100)) + batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) + + scenes = open_scenes_list() + clay = load_clay() + + for i in range(index * items_per_job, (index + 1) * items_per_job): + process_scene( + clay=clay, + path=scenes[i], + batchsize=batchsize, + ) + + +if __name__ == "__main__": + logger.debug("Starting") + process() + logger.debug("Done!") diff --git a/embeddings/environment.yml b/embeddings/environment.yml new file mode 100644 index 00000000..2c047768 --- /dev/null +++ b/embeddings/environment.yml @@ -0,0 +1,40 @@ +name: claymodel +channels: + - conda-forge + - nodefaults +dependencies: + - conda-lock~=2.5.6 + - einops~=0.7.0 + - fiona~=1.9.5 + - geopandas-base~=0.14.1 + - jsonargparse~=4.27.0 + - lightning~=2.1.0 + - matplotlib-base~=3.8.2 + - planetary-computer~=1.0.0 + - python-box~=7.1.0 + - python~=3.11.0 + - pyarrow~=16.1.0 + - rasterio~=1.3.10 + - s3fs~=2024.3.1 + - scikit-image~=0.22.0 + - scikit-learn~=1.4.0 + - stackstac~=0.5.0 + - timm~=0.9.16 + - torchvision~=0.18.1 + - transformers~=4.35.2 + - typeshed-client~=2.4.0 + - vit-pytorch~=1.6.4 + - zarr~=2.16.1 + - pip: + - geoarrow-pyarrow==0.1.2 + - jupyter-book==1.0.2 + - jupyterlab==4.2.4 + - onnx==1.16.1 + - onnxscript + - onnxruntime + - torchdata==0.7.1 + - torchgeo==0.5.2 + - wandb==0.17.5 + - stacchip==0.1.38 +platforms: + - linux-64 diff --git a/embeddings/strategy.sh b/embeddings/strategy.sh new file mode 100644 index 00000000..1d84625c --- /dev/null +++ b/embeddings/strategy.sh @@ -0,0 +1,47 @@ + # TODO: Remove for final + with open("clay_embeddings.parquet", "wb") as dst: + dst.write(body) + import geopandas as gpd + df = gpd.read_parquet("clay_embeddings.parquet") + del df["embeddings"] + df.crs = 4326 + df.to_file("clay_embeddings.gpkg") + + + + # with open("/Users/tam/Desktop/m_3008501_ne_16_1_20110815.tif") as f: + + +aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-2.amazonaws.com +docker pull 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 + + +aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 875815656045.dkr.ecr.us-east-1.amazonaws.com +docker build -t clay-v1-naip-embeddings . +docker tag clay-v1-naip-embeddings:latest 875815656045.dkr.ecr.us-east-1.amazonaws.com/clay-v1-naip-embeddings:latest +docker push 875815656045.dkr.ecr.us-east-1.amazonaws.com/clay-v1-naip-embeddings:latest + + +docker build -t clay-v1-naip-embeddings -f docker/Dockerfile . + +docker run --rm -it \ + --cpus=6 \ + -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \ + -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \ + -e AWS_BATCH_JOB_ARRAY_INDEX=0 \ + -e EMBEDDING_BATCH_SIZE=20 \ + -v /Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif:/data/m_4911964_sw_11_060_20210627.tif \ + -v $PWD/docker/all-naip.py:/code/all-naip.py \ + clay-v1-naip-embeddings + + + + item = create_stac_item("/Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif", with_proj=True)# TODO: remove! +MANIFEST = "/Users/tam/Documents/repos/model/data/naip-manifest.txt.zip" +CHECKPOINT = "/Users/tam/Documents/repos/model/data/checkpoints/clay-model-v1.5.0-september-30.ckpt" + + + + item = create_stac_item( + "/Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif", with_proj=True + ) # TODO: remove! diff --git a/embeddings/utils.py b/embeddings/utils.py new file mode 100644 index 00000000..324f4535 --- /dev/null +++ b/embeddings/utils.py @@ -0,0 +1,168 @@ +import logging +import math + +import boto3 +import geoarrow.pyarrow as ga +import numpy as np +import pyarrow as pa +import torch +from geoarrow.pyarrow import io as gaio +from stacchip.chipper import Chipper +from torchvision.transforms import v2 + +from src.module import ClayMAEModule + +CHECKPOINT = "/data/clay-model-v1.5.0-october-12.ckpt" +EMBEDDING_SHAPE_CLASS = 2 +EMBEDDING_SHAPE_PATCH = 3 + +logger = logging.getLogger("clay") + + +def normalize_timestamp(date): + week = date.isocalendar().week * 2 * np.pi / 52 + hour = date.hour * 2 * np.pi / 24 + + return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour)) + + +def normalize_latlon(lat, lon): + lat = lat * np.pi / 180 + lon = lon * np.pi / 180 + + return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon)) + + +def prepare_datacube(mean, std, datetimes, bboxs, pixels, gsd): + transform = v2.Compose( + [ + v2.Normalize(mean=mean, std=std), + ] + ) + + times = [normalize_timestamp(dat) for dat in datetimes] + week_norm = [dat[0] for dat in times] + hour_norm = [dat[1] for dat in times] + time_norm = np.hstack((week_norm, hour_norm)) + + latlons = [normalize_latlon(*bbox.centroid.coords[0]) for bbox in bboxs] + lat_norm = [dat[0] for dat in latlons] + lon_norm = [dat[1] for dat in latlons] + latlon_norm = np.hstack((lat_norm, lon_norm)) + + gsd = [gsd] + + pixels_norm = transform(pixels) + + return time_norm, latlon_norm, gsd, pixels_norm + + +def get_pixels(item, indexer_class): + indexer = indexer_class(item) + + # Instanciate the chipper + chipper = Chipper(indexer) + + # Get first chip for the "image" asset key + chips = [] + datetimes = [] + bboxs = [] + chip_ids = [] + item_ids = [] + for x, y, chip in chipper: + chips.append(chip) + datetimes.append(item.datetime) + bboxs.append(indexer.get_chip_bbox(x, y)) + chip_ids.append((x, y)) + item_ids.append(item.id) + + pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips]) + return bboxs, datetimes, pixels + + +def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchsize): # noqa: PLR0913 + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + logger.debug(f"Using device {device} to create {len(pixels_norm)} embeddings") + # Run the clay encoder + cls_embeddings = None + for i in range(0, len(pixels_norm), batchsize): + if i / batchsize % 5 == 0: + logger.debug(f"Iteration {i}") + + datacube = { + "pixels": torch.tensor( + pixels_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "time": torch.tensor( + time_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "latlon": torch.tensor( + latlon_norm[i : (i + batchsize)], dtype=torch.float32, device=device + ), + "waves": torch.tensor(waves, dtype=torch.float32, device=device), + "gsd": torch.tensor(gsd, dtype=torch.float32, device=device), + "platform": ["naip"], + } + with torch.no_grad(): + unmsk_patch, unmsk_idx, msk_idx, msk_matrix = clay.model.encoder(datacube) + # The first embedding is the class token, which is the + # overall single embedding we want to keep. + if cls_embeddings is None: + cls_embeddings = unmsk_patch[:, 0, :] + patch_embeddings = unmsk_patch[:, 1:, :] + else: + cls_embeddings = torch.vstack((cls_embeddings, unmsk_patch[:, 0, :])) + patch_embeddings = torch.vstack((patch_embeddings, unmsk_patch[:, 1:, :])) + + return cls_embeddings, patch_embeddings + + +def load_clay(): + device = "cuda" if torch.cuda.is_available() else "cpu" + model = ClayMAEModule.load_from_checkpoint( + checkpoint_path=CHECKPOINT, + metadata_path="configs/metadata.yaml", + model_size="large", + dolls=[16, 32, 64, 128, 256, 768, 1024], + doll_weights=[1, 1, 1, 1, 1, 1, 1], + mask_ratio=0.0, + shuffle=False, + ) + model.eval() + + return model.to(device) + + +def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): + np_embeddings = embeddings.cpu().numpy() + index = {"geometry": ga.as_geoarrow([dat.wkt for dat in bboxs])} + if len(embeddings.shape) == EMBEDDING_SHAPE_CLASS: + # Handle class embeddings + index["embeddings"] = [np.ascontiguousarray(dat) for dat in np_embeddings] + embedding_level = "class" + elif len(embeddings.shape) == EMBEDDING_SHAPE_PATCH: + # Handle patch embeddings + for i in range(embeddings.shape[1]): + index[f"patch_embeddings_{i}"] = [ + np.ascontiguousarray(dat) for dat in np_embeddings[:, i, :] + ] + embedding_level = "patch" + + table = pa.table( + index, + metadata={ + "date": datestr, + "gsd": str(gsd[0]), + "uri": f"s3://naip-analytic/{path}", + }, + ) + + writer = pa.BufferOutputStream() + gaio.write_geoparquet_table(table, writer) + body = bytes(writer.getvalue()) + s3_resource = boto3.resource("s3") + s3_bucket = s3_resource.Bucket(name=destination_bucket) + s3_bucket.put_object( + Body=body, + Key=f"{embedding_level}/{path.parent}/{path.stem}.parquet", + ) From c71bebf22c073ad716fe4a7d36590e0c4736f53a Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 22:55:49 +0100 Subject: [PATCH 42/83] Upgrade stacchip --- environment.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/environment.yml b/environment.yml index ac0ddeaf..fed9665d 100644 --- a/environment.yml +++ b/environment.yml @@ -36,7 +36,7 @@ dependencies: - onnxruntime - torchdata==0.7.1 - torchgeo==0.5.2 - - stacchip==0.1.35 + - stacchip==0.1.38 - wandb==0.17.5 platforms: - linux-64 From 97404e8a8854bf10b7157417666e77bd4487773b Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 22:58:29 +0100 Subject: [PATCH 43/83] Update adding sentinel-2 --- embeddings/README.md | 28 +++++++++ embeddings/parse-sentinel-2-inventory.py | 80 ++++++++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 embeddings/README.md create mode 100644 embeddings/parse-sentinel-2-inventory.py diff --git a/embeddings/README.md b/embeddings/README.md new file mode 100644 index 00000000..7ff1d6e7 --- /dev/null +++ b/embeddings/README.md @@ -0,0 +1,28 @@ +## Large scale embedding runs + +The code in this section has been used to create embedding runs over large +archives. Currently this covers NAIP and Sentinel-2. + +The algorithms are dockerized to be ran in a batch setup. AWS Batch is what +was used to execute the algorithms but it is not a strict requirement. + +The scripts rely on the `AWS_BATCH_JOB_ARRAY_INDEX` environment variable +to choose which files from the archives to process. This is set automatically +by AWS Batch when using array jobs. Outside of array jobs, this index variable +needs to be specified manually. + +### NAIP + +For NAIP, we use the `naip-analytic` bucket. We leverage the manifest file that +lists all files in the bucket. This list is parsed in the beginning and each +job processes a section of the naip scenes. + +### Sentinel-2 + +For Sentinel-2 we use the `sentinel-cogs` bucket. Also here we use the manifest +file, but parse it beforehand because it contains references to each single +asset for each product. + +The parser is essentially copied from [this gist](https://github.com/alexgleith/sinergise-element84-sentinel-2-qa/blob/main/0-parse-inventory-element84.py) +by @alexgleith. +The resulting zip file contains a list of static STAC json files for 2023 and 2024. diff --git a/embeddings/parse-sentinel-2-inventory.py b/embeddings/parse-sentinel-2-inventory.py new file mode 100644 index 00000000..397674d8 --- /dev/null +++ b/embeddings/parse-sentinel-2-inventory.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +# From https://github.com/alexgleith/sinergise-element84-sentinel-2-qa/blob/main/0-parse-inventory-element84.py +import csv +import gzip +import json +import sys + +import boto3 + +SPECIAL_YEAR = "2019" +CUTOFF_YEAR = 2023 + +s3 = boto3.resource("s3") + +bucket = "sentinel-cogs-inventory" +manifest_key = "sentinel-cogs/sentinel-cogs/2024-10-03T01-00Z/manifest.json" + +print("Starting up...") + + +def log(comment): + sys.stdout.write(f"\r{comment}") + + +# Stolen from https://alukach.com/posts/parsing-s3-inventory-output +def list_keys(bucket, manifest_key): + manifest = json.load(s3.Object(bucket, manifest_key).get()["Body"]) + for item in manifest["files"]: + gzip_obj = s3.Object(bucket_name=bucket, key=item["key"]) + buffer = gzip.open(gzip_obj.get()["Body"], mode="rt") + reader = csv.reader(buffer) + yield from reader + + +limit = 2 +count = 0 +valid = 0 +log_every = 10000 +cutoff_year = 2023 + +if __name__ == "__main__": + # Parse zip file for all scenes + with gzip.open("data/element84-tiles.list.gz", "wt") as text_file: + for tiles_bucket, key, *rest in list_keys(bucket, manifest_key): + if ".json" in key: + c = key.split("/") + # Counting scenes + count += 1 + if count % log_every == 0: + log(f"Found {count} scenes...") + tile = f"{c[1]}{c[2]}{c[3]}" + text_file.write(f"s3://{tiles_bucket}/{key}\n") + + print(f"Found {count} scenes") + + # Reduce to 2023 and 20204 + with gzip.open("data/element84-tiles-2023.gz", "wt") as dst: + with gzip.open("data/element84-tiles.list.gz") as fl: + line = fl.readline() + while line: + line = line.decode().rstrip() + c = line.split("/") + # Skip data befor 2023. Some scenes from 2019 have the year + # in a different part of the prefix. + if c[4] == SPECIAL_YEAR: + line = fl.readline() + continue + elif int(c[7]) < CUTOFF_YEAR: + line = fl.readline() + continue + elif not line.endswith("L2A.json"): + line = fl.readline() + continue + + count += 1 + if count % log_every == 0: + log(f"Found {count} scenes... {line}") + + dst.write(line + "\n") + line = fl.readline() From 308eb8010ef4e0010b957e99e8922c8ee6403bce Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 23:05:43 +0100 Subject: [PATCH 44/83] Update docker file --- embeddings/Dockerfile | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 3e32308e..c3f48259 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -3,16 +3,11 @@ FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt /data/clay-model-v1.5.0-october-12.ckpt +RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip /data/naip-manifest.txt.zip RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model # ADD docker/environment.yml docker-environment.yml RUN mamba env create --file environment.yml -ADD data/naip-manifest.txt.zip /data/naip-manifest.txt.zip -ADD data/checkpoints/clay-model-v1.5.0-september-30.ckpt /data/clay-model-v1.5.0-september-30.ckpt - -ADD docker/all-naip.py . -ADD configs configs - -ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "docker/all-naip.py"] From c9af8e64f2c9738bc8bd43e545e99a72dbef9fde Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 23:07:43 +0100 Subject: [PATCH 45/83] Update docker file --- embeddings/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index c3f48259..be279da8 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -1,4 +1,4 @@ -FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 +FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 WORKDIR /model From 8ce1515f9b4940f55449b35daba9a1db69956ec6 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 23:12:54 +0100 Subject: [PATCH 46/83] Update docker file --- embeddings/Dockerfile | 2 +- embeddings/README.md | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index be279da8..018fc1d3 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -10,4 +10,4 @@ RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model # ADD docker/environment.yml docker-environment.yml RUN mamba env create --file environment.yml -ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "docker/all-naip.py"] +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "embeddings/all-naip.py"] diff --git a/embeddings/README.md b/embeddings/README.md index 7ff1d6e7..c0333893 100644 --- a/embeddings/README.md +++ b/embeddings/README.md @@ -11,6 +11,19 @@ to choose which files from the archives to process. This is set automatically by AWS Batch when using array jobs. Outside of array jobs, this index variable needs to be specified manually. +### Build docker image + +Build the docker image from the embeddings directory and push to ECR +or another docker repository of your choice. + +```bash +aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-2.amazonaws.com +docker pull 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 +docker build -t clay-embeddings -f embeddings/Dockerfile . +docker tag clay-embeddings:latest 875815656045.dkr.ecr.us-east-2.amazonaws.com/clay-embeddings:latest +docker push 875815656045.dkr.ecr.us-east-2.amazonaws.com/clay-embeddings:latest +``` + ### NAIP For NAIP, we use the `naip-analytic` bucket. We leverage the manifest file that From 300d978f36e16915d2c9745dfe10f1d6605482e0 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Mon, 14 Oct 2024 23:31:40 +0100 Subject: [PATCH 47/83] update dockerfile --- embeddings/Dockerfile | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 018fc1d3..f9d4b5b9 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -7,7 +7,6 @@ RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip /data/naip-manifest.t RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model -# ADD docker/environment.yml docker-environment.yml -RUN mamba env create --file environment.yml +RUN mamba env create --file embeddings/environment.yml ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "embeddings/all-naip.py"] From 9b5b7a6329c05a5b69775f7c0cfc5261d3422537 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 11:13:51 +0100 Subject: [PATCH 48/83] Fix import path --- embeddings/README.md | 8 ++++++-- embeddings/all-naip.py | 4 ++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/embeddings/README.md b/embeddings/README.md index c0333893..302b0278 100644 --- a/embeddings/README.md +++ b/embeddings/README.md @@ -13,14 +13,18 @@ needs to be specified manually. ### Build docker image -Build the docker image from the embeddings directory and push to ECR -or another docker repository of your choice. +Embedding runs are dockerized for parallel computing. To build the docker image +use the Dockerfile in the embeddings directory. Then push the image to ECR or +another docker repository of your choice. ```bash aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-2.amazonaws.com docker pull 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 + docker build -t clay-embeddings -f embeddings/Dockerfile . + docker tag clay-embeddings:latest 875815656045.dkr.ecr.us-east-2.amazonaws.com/clay-embeddings:latest +aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin 875815656045.dkr.ecr.us-east-2.amazonaws.com docker push 875815656045.dkr.ecr.us-east-2.amazonaws.com/clay-embeddings:latest ``` diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 044a6b5c..79ee493b 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -11,7 +11,7 @@ from rio_stac import create_stac_item from stacchip.indexer import NoStatsChipIndexer -from embeddings.utils import ( +from utils import ( get_embeddings, get_pixels, load_clay, @@ -112,7 +112,7 @@ def process(): raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) items_per_job = int(os.environ.get("ITEMS_PER_JOB", 2)) - batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 2)) + batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) scenes = open_scene_list() clay = load_clay() From d79f93e60b86d174fa6b1e209686bc889e672453 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 13:37:51 +0100 Subject: [PATCH 49/83] Fix paths again --- embeddings/Dockerfile | 6 +++++- embeddings/all-naip.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index f9d4b5b9..557d10c6 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -9,4 +9,8 @@ RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model RUN mamba env create --file embeddings/environment.yml -ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "embeddings/all-naip.py"] +# Move file to home directory so that relative imports work +RUN cp embeddings/all-naip.py . +RUN cp embeddings/all-sentinel.py . + +ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 79ee493b..9330616c 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -11,7 +11,7 @@ from rio_stac import create_stac_item from stacchip.indexer import NoStatsChipIndexer -from utils import ( +from embeddings.utils import ( get_embeddings, get_pixels, load_clay, From 4d8a1a5926b8ad12d3478614d1b692622b448738 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 14:05:11 +0100 Subject: [PATCH 50/83] Fix paths in docker file --- embeddings/Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 557d10c6..3f3fdcdc 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -2,8 +2,8 @@ FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model -RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt /data/clay-model-v1.5.0-october-12.ckpt -RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip /data/naip-manifest.txt.zip +RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt +RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model From 70030ba905d2928949016092db42c452873e30c4 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 14:23:46 +0100 Subject: [PATCH 51/83] Fix docker file --- embeddings/Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 3f3fdcdc..16480f8c 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -2,11 +2,11 @@ FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model +RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model + RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip -RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model - RUN mamba env create --file embeddings/environment.yml # Move file to home directory so that relative imports work From 7beb600e0eb672467f9c70e841ffba03ce8412d0 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 14:23:54 +0100 Subject: [PATCH 52/83] Fix docker file --- embeddings/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 16480f8c..b8e448c4 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -2,7 +2,7 @@ FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model -RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git /model +RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip From d4426617e8a9570499ece3ed03405bbaf6f7807b Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 15:52:42 +0100 Subject: [PATCH 53/83] Use relative path --- embeddings/all-naip.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 9330616c..8628379c 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -24,8 +24,7 @@ logger.setLevel(logging.DEBUG) -# MANIFEST = "/data/naip-manifest.txt.zip" -MANIFEST = "/Users/tam/Documents/repos/model/data/naip-manifest.txt.zip" +MANIFEST = "data/naip-manifest.txt.zip" MEAN = [ 110.16, 115.41, From dd0d0f9e66c039dc3781fd224d273fb134ffdbff Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 15:56:48 +0100 Subject: [PATCH 54/83] Make checkpoint path relative --- embeddings/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/utils.py b/embeddings/utils.py index 324f4535..afcfbcee 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -12,7 +12,7 @@ from src.module import ClayMAEModule -CHECKPOINT = "/data/clay-model-v1.5.0-october-12.ckpt" +CHECKPOINT = "data/clay-model-v1.5.0-october-12.ckpt" EMBEDDING_SHAPE_CLASS = 2 EMBEDDING_SHAPE_PATCH = 3 From 40659da8e673e12bc126ff8d0d4d22b04acc1a31 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 16:00:20 +0100 Subject: [PATCH 55/83] Make path relative --- embeddings/all-sentinel.py | 69 +++++++------------------------------- embeddings/strategy.sh | 47 -------------------------- 2 files changed, 12 insertions(+), 104 deletions(-) delete mode 100644 embeddings/strategy.sh diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index e9c6fc1a..5f1bb348 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -2,7 +2,6 @@ import json import logging import os -import sys import boto3 from pystac import Item @@ -13,71 +12,27 @@ logger = logging.getLogger("clay") -# Create 2023 file from full archive - -s3 = boto3.resource("s3") - - -def log(comment): - sys.stdout.write(f"\r{comment}") - - -log_every = 100000 - -count = 0 - -# with gzip.open(f"data/element84-tiles-2023.gz", "wt") as dst: -# with gzip.open(f"data/element84-tiles.list.gz") as fl: -# line = fl.readline() -# while line: -# line = line.decode().rstrip() -# c = line.split("/") -# if c[4] == "2019": -# line = fl.readline() -# continue -# elif int(c[7]) < 2023: -# line = fl.readline() -# continue -# elif not line.endswith("L2A.json"): -# line = fl.readline() -# continue - -# count += 1 -# if count % log_every == 0: -# log(f"Found {count} scenes... {line}") - -# dst.write(line + "\n") -# line = fl.readline() - - -def open_manifest(path): - pass - - -data = open_manifest() - -index = 42 - -key = data[index].decode().rstrip().replace("s3://sentinel-cogs/", "") - -stac_json = json.load(s3.Object("sentinel-cogs", key).get()["Body"]) +SCENES_LIST = "data/element84-tiles-2023.gz" -item = Item.from_dict(stac_json) +def open_scenes_list(): + with gzip.open(SCENES_LIST) as fl: + data = fl.readlines() + return [dat.decode().rstrip() for dat in data] -indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) +def process_scene(clay, path, batchsize): + key = path.replace("s3://sentinel-cogs/", "") -SCENES_LIST = "data/element84-tiles-2023.gz" + s3 = boto3.resource("s3") + stac_json = json.load(s3.Object("sentinel-cogs", key).get()["Body"]) -def open_scenes_list(): - with gzip.open(SCENES_LIST) as fl: - return fl.readlines() + item = Item.from_dict(stac_json) + indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) -def process_scene(clay, path, batchsize): - pass + return indexer def process(): diff --git a/embeddings/strategy.sh b/embeddings/strategy.sh deleted file mode 100644 index 1d84625c..00000000 --- a/embeddings/strategy.sh +++ /dev/null @@ -1,47 +0,0 @@ - # TODO: Remove for final - with open("clay_embeddings.parquet", "wb") as dst: - dst.write(body) - import geopandas as gpd - df = gpd.read_parquet("clay_embeddings.parquet") - del df["embeddings"] - df.crs = 4326 - df.to_file("clay_embeddings.gpkg") - - - - # with open("/Users/tam/Desktop/m_3008501_ne_16_1_20110815.tif") as f: - - -aws ecr get-login-password --region us-east-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-2.amazonaws.com -docker pull 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 - - -aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 875815656045.dkr.ecr.us-east-1.amazonaws.com -docker build -t clay-v1-naip-embeddings . -docker tag clay-v1-naip-embeddings:latest 875815656045.dkr.ecr.us-east-1.amazonaws.com/clay-v1-naip-embeddings:latest -docker push 875815656045.dkr.ecr.us-east-1.amazonaws.com/clay-v1-naip-embeddings:latest - - -docker build -t clay-v1-naip-embeddings -f docker/Dockerfile . - -docker run --rm -it \ - --cpus=6 \ - -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \ - -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \ - -e AWS_BATCH_JOB_ARRAY_INDEX=0 \ - -e EMBEDDING_BATCH_SIZE=20 \ - -v /Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif:/data/m_4911964_sw_11_060_20210627.tif \ - -v $PWD/docker/all-naip.py:/code/all-naip.py \ - clay-v1-naip-embeddings - - - - item = create_stac_item("/Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif", with_proj=True)# TODO: remove! -MANIFEST = "/Users/tam/Documents/repos/model/data/naip-manifest.txt.zip" -CHECKPOINT = "/Users/tam/Documents/repos/model/data/checkpoints/clay-model-v1.5.0-september-30.ckpt" - - - - item = create_stac_item( - "/Users/tam/Desktop/m_4911964_sw_11_060_20210627.tif", with_proj=True - ) # TODO: remove! From 2b1a567ebe86536a35888529a4239c8a80fd508c Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 22:46:10 +0100 Subject: [PATCH 56/83] Make metadata dynamic --- embeddings/all-naip.py | 23 ++++++++--------- embeddings/all-sentinel.py | 51 +++++++++++++++++++++++++++++++++++--- embeddings/utils.py | 23 +++++++++++------ 3 files changed, 74 insertions(+), 23 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 8628379c..06e28d9d 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -9,12 +9,14 @@ import boto3 from rasterio.errors import RasterioIOError from rio_stac import create_stac_item +from stacchip.chipper import Chipper from stacchip.indexer import NoStatsChipIndexer from embeddings.utils import ( get_embeddings, get_pixels, load_clay, + load_metadata, prepare_datacube, write_to_table, ) @@ -25,14 +27,6 @@ MANIFEST = "data/naip-manifest.txt.zip" -MEAN = [ - 110.16, - 115.41, - 98.15, - 139.04, -] -STD = [47.23, 39.82, 35.43, 49.86] -WAVES = [0.65, 0.56, 0.48, 0.842] EMBEDDINGS_BUCKET = "clay-embeddings-naip" @@ -60,6 +54,8 @@ def process_scene(clay, path, batchsize): datestr = path.stem.split("_")[-1] date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) gsd = float(path.parts[2].replace("cm", "")) / 100 + bands, waves, mean, std = load_metadata("naip") + logger.debug(f"Processing {path} in state {state} and date {date}") with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as fl: @@ -74,13 +70,17 @@ def process_scene(clay, path, batchsize): item.id = f"{state}_{path.stem}" try: - bboxs, datetimes, pixels = get_pixels(item, NoStatsChipIndexer) + indexer = NoStatsChipIndexer(item) + chipper = Chipper(indexer, assets=bands) + bboxs, datetimes, pixels = get_pixels( + item=item, indexer=indexer, chipper=chipper + ) except RasterioIOError: logger.warning("Skipping scene due to rasterio io error") return time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( - mean=MEAN, std=STD, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd + mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd ) # Embed data cls_embeddings, patch_embeddings = get_embeddings( @@ -88,12 +88,11 @@ def process_scene(clay, path, batchsize): pixels_norm=pixels_norm, time_norm=time_norm, latlon_norm=latlon_norm, - waves=WAVES, + waves=waves, gsd=gsd, batchsize=batchsize, ) # Write class embeddings - kwargs = dict( bboxs=bboxs, datestr=datestr, diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 5f1bb348..f9b95bcb 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -5,14 +5,25 @@ import boto3 from pystac import Item +from rasterio.errors import RasterioIOError +from stacchip.chipper import Chipper from stacchip.indexer import Sentinel2Indexer -from embeddings.utils import load_clay +from embeddings.utils import ( + get_embeddings, + get_pixels, + load_clay, + load_metadata, + prepare_datacube, + write_to_table, +) logger = logging.getLogger("clay") SCENES_LIST = "data/element84-tiles-2023.gz" +EMBEDDINGS_BUCKET = "clay-embeddings-sentinel-2" +GSD = 10 def open_scenes_list(): @@ -22,6 +33,8 @@ def open_scenes_list(): def process_scene(clay, path, batchsize): + bands, waves, mean, std = load_metadata("sentinel_2_l2a") + key = path.replace("s3://sentinel-cogs/", "") s3 = boto3.resource("s3") @@ -30,9 +43,41 @@ def process_scene(clay, path, batchsize): item = Item.from_dict(stac_json) - indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) + bands, waves, mean, std = load_metadata("naip") - return indexer + try: + indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) + chipper = Chipper(item, assets=bands) + bboxs, datetimes, pixels = get_pixels( + item=item, indexer=indexer, chipper=chipper + ) + except RasterioIOError: + logger.warning("Skipping scene due to rasterio io error") + return + + time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( + mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=GSD + ) + # Embed data + cls_embeddings, patch_embeddings = get_embeddings( + clay=clay, + pixels_norm=pixels_norm, + time_norm=time_norm, + latlon_norm=latlon_norm, + waves=waves, + gsd=gsd, + batchsize=batchsize, + ) + kwargs = dict( + bboxs=bboxs, + datestr=str(item.datetime.date()), + gsd=gsd, + destination_bucket=EMBEDDINGS_BUCKET, + path=path, + ) + + write_to_table(embeddings=cls_embeddings, **kwargs) + write_to_table(embeddings=patch_embeddings, **kwargs) def process(): diff --git a/embeddings/utils.py b/embeddings/utils.py index afcfbcee..8e569fea 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -6,8 +6,9 @@ import numpy as np import pyarrow as pa import torch +import yaml +from box import Box from geoarrow.pyarrow import io as gaio -from stacchip.chipper import Chipper from torchvision.transforms import v2 from src.module import ClayMAEModule @@ -19,6 +20,18 @@ logger = logging.getLogger("clay") +def load_metadata(platform): + metadata = Box(yaml.safe_load(open("configs/metadata.yaml"))) + platform_meta = getattr(metadata, platform) + + bands = list(platform_meta.bands.wavelength.keys()) + waves = list(platform_meta.bands.wavelength.values()) + mean = list(platform_meta.bands.mean.values()) + std = list(platform_meta.bands.std.values()) + + return bands, waves, mean, std + + def normalize_timestamp(date): week = date.isocalendar().week * 2 * np.pi / 52 hour = date.hour * 2 * np.pi / 24 @@ -57,13 +70,7 @@ def prepare_datacube(mean, std, datetimes, bboxs, pixels, gsd): return time_norm, latlon_norm, gsd, pixels_norm -def get_pixels(item, indexer_class): - indexer = indexer_class(item) - - # Instanciate the chipper - chipper = Chipper(indexer) - - # Get first chip for the "image" asset key +def get_pixels(item, indexer, chipper): chips = [] datetimes = [] bboxs = [] From 3ad580caeca5a3057e9934b6240c85e41823f350 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 23:29:44 +0100 Subject: [PATCH 57/83] Make metadata dynamic fix --- embeddings/all-naip.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 06e28d9d..3a067940 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -71,7 +71,7 @@ def process_scene(clay, path, batchsize): try: indexer = NoStatsChipIndexer(item) - chipper = Chipper(indexer, assets=bands) + chipper = Chipper(indexer) bboxs, datetimes, pixels = get_pixels( item=item, indexer=indexer, chipper=chipper ) From c08926ee5340d85100a73fc73859a3fba6068174 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 15 Oct 2024 23:31:33 +0100 Subject: [PATCH 58/83] Log device early --- embeddings/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/embeddings/utils.py b/embeddings/utils.py index 8e569fea..3dff287f 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -126,6 +126,7 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs def load_clay(): device = "cuda" if torch.cuda.is_available() else "cpu" + logger.debug(f"Loading model on device {device}") model = ClayMAEModule.load_from_checkpoint( checkpoint_path=CHECKPOINT, metadata_path="configs/metadata.yaml", From c14f8a17e56532528c4094f49c5a8d27cf49b147 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 10:04:22 +0100 Subject: [PATCH 59/83] Use pip instead of conda to avoid re-install of torch --- embeddings/Dockerfile | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index b8e448c4..42f1872b 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -2,15 +2,38 @@ FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model -RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . - RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip -RUN mamba env create --file embeddings/environment.yml +RUN pip install \ + einops~=0.7.0 \ + fiona~=1.9.5 \ + geopandas~=0.14.1 \ + jsonargparse~=4.27.0 \ + lightning~=2.1.0 \ + matplotlib~=3.8.2 \ + planetary-computer~=1.0.0 \ + python-box~=7.1.0 \ + pyarrow~=16.1.0 \ + rasterio~=1.3.10 \ + s3fs~=2024.3.1 \ + scikit-image~=0.22.0 \ + scikit-learn~=1.4.0 \ + stackstac~=0.5.0 \ + timm~=0.9.16 \ + transformers~=4.35.2 \ + typeshed-client~=2.4.0 \ + vit-pytorch~=1.6.4 \ + torch~=2.3.0 \ + zarr~=2.16.1 \ + geoarrow-pyarrow==0.1.2 \ + torchdata==0.7.1 \ + stacchip==0.1.35 \ + wandb==0.17.5 +RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . RUN cp embeddings/all-sentinel.py . -ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] +ENTRYPOINT ["python", "all-naip.py"] From 9952c4efaa8d173be2b50a31f51aecb67d16abfa Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 10:05:41 +0100 Subject: [PATCH 60/83] Use pip instead of conda to avoid re-install of torch --- embeddings/Dockerfile | 1 - 1 file changed, 1 deletion(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 42f1872b..e2289f92 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -24,7 +24,6 @@ RUN pip install \ transformers~=4.35.2 \ typeshed-client~=2.4.0 \ vit-pytorch~=1.6.4 \ - torch~=2.3.0 \ zarr~=2.16.1 \ geoarrow-pyarrow==0.1.2 \ torchdata==0.7.1 \ From f2773eed8957438a2f8f4cf004a548f4c5b700d7 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 10:23:51 +0100 Subject: [PATCH 61/83] Use pip instead of conda to avoid re-install of torch --- embeddings/Dockerfile | 8 +++++--- embeddings/all-naip.py | 3 ++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index e2289f92..c4e10ede 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -2,8 +2,7 @@ FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-inference:2.3.0-gpu-py WORKDIR /model -RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt -RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip +RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . RUN pip install \ einops~=0.7.0 \ @@ -30,7 +29,10 @@ RUN pip install \ stacchip==0.1.35 \ wandb==0.17.5 -RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . + +RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt +RUN aws s3 cp s3://clay-mgrs-samples/naipgi-manifest.txt.zip data/naip-manifest.txt.zip + # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . RUN cp embeddings/all-sentinel.py . diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 3a067940..9e6e65d4 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -100,8 +100,9 @@ def process_scene(clay, path, batchsize): destination_bucket=EMBEDDINGS_BUCKET, path=path, ) + logger.debug("Writing class embeddings") write_to_table(embeddings=cls_embeddings, **kwargs) - # Write patch embeddings + logger.debug("Writing patch embeddings") write_to_table(embeddings=patch_embeddings, **kwargs) From cdf6e41f892c4ab239b3fbe85e2b3a2b38e13909 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 10:30:07 +0100 Subject: [PATCH 62/83] Use pip instead of conda to avoid re-install of torch --- embeddings/Dockerfile | 2 ++ 1 file changed, 2 insertions(+) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index c4e10ede..05681f07 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -16,6 +16,8 @@ RUN pip install \ pyarrow~=16.1.0 \ rasterio~=1.3.10 \ s3fs~=2024.3.1 \ + boto3~=1.34.122 \ + botocore~=1.34.122 \ scikit-image~=0.22.0 \ scikit-learn~=1.4.0 \ stackstac~=0.5.0 \ From bdb2ca3c4097600e95b4df04c34dfc9679824194 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 10:36:28 +0100 Subject: [PATCH 63/83] Use pip instead of conda to avoid re-install of torch --- embeddings/Dockerfile | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 05681f07..62037461 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -10,12 +10,12 @@ RUN pip install \ geopandas~=0.14.1 \ jsonargparse~=4.27.0 \ lightning~=2.1.0 \ - matplotlib~=3.8.2 \ + matplotlib~=3.9.0 \ planetary-computer~=1.0.0 \ python-box~=7.1.0 \ - pyarrow~=16.1.0 \ + pyarrow~=15.0.2 \ rasterio~=1.3.10 \ - s3fs~=2024.3.1 \ + s3fs~=2024.6.0 \ boto3~=1.34.122 \ botocore~=1.34.122 \ scikit-image~=0.22.0 \ @@ -33,7 +33,7 @@ RUN pip install \ RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt -RUN aws s3 cp s3://clay-mgrs-samples/naipgi-manifest.txt.zip data/naip-manifest.txt.zip +RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . From 9146bddcaa0e09e5c2dc0c65a512631e32fa023e Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 14:30:49 +0100 Subject: [PATCH 64/83] Remove patch level embeddings storing --- embeddings/all-naip.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 9e6e65d4..bf8193ab 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -102,8 +102,6 @@ def process_scene(clay, path, batchsize): ) logger.debug("Writing class embeddings") write_to_table(embeddings=cls_embeddings, **kwargs) - logger.debug("Writing patch embeddings") - write_to_table(embeddings=patch_embeddings, **kwargs) def process(): From b093ec526d447c61f0f863f6f7017fdd44dd8ac1 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 16 Oct 2024 18:47:38 +0100 Subject: [PATCH 65/83] Flaten path to mirror naip-analytic bucket --- embeddings/README.md | 2 ++ embeddings/utils.py | 4 +--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/embeddings/README.md b/embeddings/README.md index 302b0278..cc7685ee 100644 --- a/embeddings/README.md +++ b/embeddings/README.md @@ -34,6 +34,8 @@ For NAIP, we use the `naip-analytic` bucket. We leverage the manifest file that lists all files in the bucket. This list is parsed in the beginning and each job processes a section of the naip scenes. +At the moment of processing there were 1'231'441 NAIP scenes. + ### Sentinel-2 For Sentinel-2 we use the `sentinel-cogs` bucket. Also here we use the manifest diff --git a/embeddings/utils.py b/embeddings/utils.py index 3dff287f..2c33dfb4 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -147,14 +147,12 @@ def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): if len(embeddings.shape) == EMBEDDING_SHAPE_CLASS: # Handle class embeddings index["embeddings"] = [np.ascontiguousarray(dat) for dat in np_embeddings] - embedding_level = "class" elif len(embeddings.shape) == EMBEDDING_SHAPE_PATCH: # Handle patch embeddings for i in range(embeddings.shape[1]): index[f"patch_embeddings_{i}"] = [ np.ascontiguousarray(dat) for dat in np_embeddings[:, i, :] ] - embedding_level = "patch" table = pa.table( index, @@ -172,5 +170,5 @@ def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): s3_bucket = s3_resource.Bucket(name=destination_bucket) s3_bucket.put_object( Body=body, - Key=f"{embedding_level}/{path.parent}/{path.stem}.parquet", + Key=f"{path.parent}/{path.stem}.parquet", ) From ab703ce5c03136e0f7b8f266d36f2cc087e5ddbd Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Thu, 17 Oct 2024 12:29:13 +0100 Subject: [PATCH 66/83] Add option to limit to state --- embeddings/all-naip.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index bf8193ab..db5dbd21 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -30,7 +30,7 @@ EMBEDDINGS_BUCKET = "clay-embeddings-naip" -def open_scene_list(): +def open_scene_list(limit_to_state=None): """ Read the naip-analytic manifest file and extract a list of NAIP scenes as tif files to process. @@ -42,7 +42,13 @@ def open_scene_list(): data = f.readlines() data = [Path(dat.rstrip()) for dat in data if "rgbir_cog"] data = [dat for dat in data if dat.suffix == ".tif"] + logger.debug(f"Found {len(data)} NAIP scenes in manifest") + + if limit_to_state is not None: + data = [dat for dat in data if str(dat).startswith(limit_to_state)] + logger.debug(f"Found {len(data)} NAIP scenes for state {limit_to_state}") + return data @@ -110,8 +116,9 @@ def process(): index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) items_per_job = int(os.environ.get("ITEMS_PER_JOB", 2)) batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) + limit_to_state = os.environ.get("LIMIT_TO_STATE", None) - scenes = open_scene_list() + scenes = open_scene_list(limit_to_state) clay = load_clay() for i in range(index * items_per_job, (index + 1) * items_per_job): From e902414d3466d2ca1dad477671932c0018b49925 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 5 Nov 2024 09:19:00 +0000 Subject: [PATCH 67/83] Remove stale files --- docker/Dockerfile | 16 --- docker/all-naip.py | 261 --------------------------------------------- 2 files changed, 277 deletions(-) delete mode 100644 docker/Dockerfile delete mode 100644 docker/all-naip.py diff --git a/docker/Dockerfile b/docker/Dockerfile deleted file mode 100644 index 9fa8d4a2..00000000 --- a/docker/Dockerfile +++ /dev/null @@ -1,16 +0,0 @@ -FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:2.3.0-gpu-py311-cu121-ubuntu20.04-ec2 - -WORKDIR /code - -ADD src src - -ADD docker/environment.yml docker-environment.yml -RUN mamba env create --file docker-environment.yml - -ADD data/naip-manifest.txt.zip /data/naip-manifest.txt.zip -ADD data/checkpoints/clay-model-v1.5.0-september-30.ckpt /data/clay-model-v1.5.0-september-30.ckpt - -ADD docker/all-naip.py . -ADD configs configs - -ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "claymodel", "python", "all-naip.py"] diff --git a/docker/all-naip.py b/docker/all-naip.py deleted file mode 100644 index bd1eddbe..00000000 --- a/docker/all-naip.py +++ /dev/null @@ -1,261 +0,0 @@ -import datetime -import io -import logging -import math -import os -import tempfile -import zipfile -from pathlib import Path - -import boto3 -import geoarrow.pyarrow as ga -import numpy as np -import pyarrow as pa -import torch -from geoarrow.pyarrow import io as gaio -from rasterio.errors import RasterioIOError -from rio_stac import create_stac_item -from stacchip.chipper import Chipper -from stacchip.indexer import NoStatsChipIndexer -from torchvision.transforms import v2 - -from src.module import ClayMAEModule - -MANIFEST = "/data/naip-manifest.txt.zip" -CHECKPOINT = "/data/clay-model-v1.5.0-september-30.ckpt" - -logging.basicConfig() -logger = logging.getLogger("clay") -logger.setLevel(logging.DEBUG) - - -def normalize_timestamp(date): - week = date.isocalendar().week * 2 * np.pi / 52 - hour = date.hour * 2 * np.pi / 24 - - return (math.sin(week), math.cos(week)), (math.sin(hour), math.cos(hour)) - - -def normalize_latlon(lat, lon): - lat = lat * np.pi / 180 - lon = lon * np.pi / 180 - - return (math.sin(lat), math.cos(lat)), (math.sin(lon), math.cos(lon)) - - -def prepare_datacube(datetimes, bboxs, pixels, gsd): - # Set mean, std, and wavelengths metadata - mean = [ - 110.16, - 115.41, - 98.15, - 139.04, - ] - std = [47.23, 39.82, 35.43, 49.86] - waves = [0.65, 0.56, 0.48, 0.842] - - transform = v2.Compose( - [ - v2.Normalize(mean=mean, std=std), - ] - ) - - times = [normalize_timestamp(dat) for dat in datetimes] - week_norm = [dat[0] for dat in times] - hour_norm = [dat[1] for dat in times] - time_norm = np.hstack((week_norm, hour_norm)) - - latlons = [normalize_latlon(*bbox.centroid.coords[0]) for bbox in bboxs] - lat_norm = [dat[0] for dat in latlons] - lon_norm = [dat[1] for dat in latlons] - latlon_norm = np.hstack((lat_norm, lon_norm)) - - gsd = [gsd] - - pixels_norm = transform(pixels) - - return waves, time_norm, latlon_norm, gsd, pixels_norm - - -def get_pixels(item): - indexer = NoStatsChipIndexer(item) - - # Instanciate the chipper - chipper = Chipper(indexer) - - # Get first chip for the "image" asset key - chips = [] - datetimes = [] - bboxs = [] - chip_ids = [] - item_ids = [] - for idx, (x, y, chip) in enumerate(chipper): - chips.append(chip) - datetimes.append(item.datetime) - bboxs.append(indexer.get_chip_bbox(x, y)) - chip_ids.append((x, y)) - item_ids.append(item.id) - - pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips]) - return bboxs, datetimes, pixels - - -def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchsize): # noqa: PLR0913 - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") - logger.debug(f"Using device {device} to create {len(pixels_norm)} embeddings") - # Run the clay encoder - embeddings = None - for i in range(0, len(pixels_norm), batchsize): - if i / batchsize % 5 == 0: - logger.debug(f"Iteration {i}") - - datacube = { - "pixels": torch.tensor( - pixels_norm[i : (i + batchsize)], dtype=torch.float32, device=device - ), - "time": torch.tensor( - time_norm[i : (i + batchsize)], dtype=torch.float32, device=device - ), - "latlon": torch.tensor( - latlon_norm[i : (i + batchsize)], dtype=torch.float32, device=device - ), - "waves": torch.tensor(waves, dtype=torch.float32, device=device), - "gsd": torch.tensor(gsd, dtype=torch.float32, device=device), - "platform": ["naip"], - } - with torch.no_grad(): - unmsk_patch, unmsk_idx, msk_idx, msk_matrix = clay.model.encoder(datacube) - # The first embedding is the class token, which is the - # overall single embedding we want to keep. - cls_embeddings = unmsk_patch[:, 0, :] - if embeddings is None: - embeddings = cls_embeddings - else: - embeddings = torch.vstack((embeddings, cls_embeddings)) - - return embeddings - - -def open_scene_list(): - with zipfile.ZipFile(MANIFEST) as zf: - with io.TextIOWrapper(zf.open("naip-manifest.txt"), encoding="utf-8") as f: - data = f.readlines() - data = [Path(dat.rstrip()) for dat in data if "rgbir_cog"] - data = [dat for dat in data if dat.suffix == ".tif"] - logger.debug(f"Found {len(data)} NAIP scenes in manifest") - return data - - -def load_clay(): - device = "cuda" if torch.cuda.is_available() else "cpu" - model = ClayMAEModule.load_from_checkpoint( - checkpoint_path=CHECKPOINT, - metadata_path="configs/metadata.yaml", - model_size="large", - dolls=[16, 32, 64, 128, 256, 768, 1024], - doll_weights=[1, 1, 1, 1, 1, 1, 1], - mask_ratio=0.0, - shuffle=False, - ) - model.eval() - - return model.to(device) - - -def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): - index = { - "embeddings": [np.ascontiguousarray(dat) for dat in embeddings.cpu().numpy()], - "geometry": ga.as_geoarrow([dat.wkt for dat in bboxs]), - } - - table = pa.table( - index, - metadata={ - "date": datestr, - "gsd": str(gsd[0]), - "uri": f"s3://naip-analytic/{path}", - }, - ) - - writer = pa.BufferOutputStream() - gaio.write_geoparquet_table(table, writer) - body = bytes(writer.getvalue()) - s3_resource = boto3.resource("s3") - s3_bucket = s3_resource.Bucket(name=destination_bucket) - s3_bucket.put_object( - Body=body, - Key=f"{path.parent}/{path.stem}.parquet", - ) - - -def process_scene(clay, path, destination_bucket, batchsize): - state = path.parts[0] - datestr = path.stem.split("_")[-1] - date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) - gsd = float(path.parts[2].replace("cm", "")) / 100 - logger.debug(f"Processing {path} in state {state} and date {date}") - - with tempfile.NamedTemporaryFile(mode="w+b", suffix=".tif") as fl: - s3 = boto3.client("s3") - s3.download_fileobj( - "naip-analytic", str(path), fl, ExtraArgs={"RequestPayer": "requester"} - ) - - item = create_stac_item(fl.name, with_proj=True) - item.datetime = date - item.id = f"{state}_{path.stem}" - - try: - bboxs, datetimes, pixels = get_pixels(item) - except RasterioIOError: - logger.warning("Skipping scene due to rasterio io error") - return - - waves, time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( - datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd - ) - - embeddings = get_embeddings( - clay=clay, - pixels_norm=pixels_norm, - time_norm=time_norm, - latlon_norm=latlon_norm, - waves=waves, - gsd=gsd, - batchsize=batchsize, - ) - - write_to_table( - embeddings=embeddings, - bboxs=bboxs, - datestr=datestr, - gsd=gsd, - destination_bucket=destination_bucket, - path=path, - ) - - -def process(): - if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: - raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") - index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) - items_per_job = int(os.environ.get("ITEMS_PER_JOB", 100)) - batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) - destination_bucket = "clay-v1-naip-embeddings" - - scenes = open_scene_list() - clay = load_clay() - - for i in range(index * items_per_job, (index + 1) * items_per_job): - process_scene( - clay=clay, - path=scenes[i], - destination_bucket=destination_bucket, - batchsize=batchsize, - ) - - -if __name__ == "__main__": - logger.debug("Starting") - process() - logger.debug("Done!") From 2f33b445a2d8f9e8028df76cfc5ea6366893d0f9 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 5 Nov 2024 09:57:17 +0000 Subject: [PATCH 68/83] Update model checkpoint name --- embeddings/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/utils.py b/embeddings/utils.py index 2c33dfb4..fba30304 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -13,7 +13,7 @@ from src.module import ClayMAEModule -CHECKPOINT = "data/clay-model-v1.5.0-october-12.ckpt" +CHECKPOINT = "data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt" EMBEDDING_SHAPE_CLASS = 2 EMBEDDING_SHAPE_PATCH = 3 From d85a664bcf0118c48a896c09627472065de8a5da Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 5 Nov 2024 09:57:47 +0000 Subject: [PATCH 69/83] Update model checkpoint and s3 sign strategy --- embeddings/Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 62037461..a8a5d7a3 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -32,8 +32,8 @@ RUN pip install \ wandb==0.17.5 -RUN aws s3 cp s3://clay-model-ckpt/v1.5.0/mae_v1.5.0_epoch-76_val-loss-0.1612.ckpt data/clay-model-v1.5.0-october-12.ckpt -RUN aws s3 cp s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip +RUN aws s3 cp --no-sign-request s3://clay-model-ckpt/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt +RUN aws s3 cp --no-sign-request s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . From 85f054448ef9fd44617e555082c1cbbbb2d480dd Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 5 Nov 2024 17:18:49 +0000 Subject: [PATCH 70/83] Adapt to rio-stac 0.10.0 and pin requirement --- embeddings/Dockerfile | 3 ++- embeddings/all-naip.py | 10 ++++++---- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index a8a5d7a3..26024c27 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -29,7 +29,8 @@ RUN pip install \ geoarrow-pyarrow==0.1.2 \ torchdata==0.7.1 \ stacchip==0.1.35 \ - wandb==0.17.5 + wandb==0.17.5 \ + rio_stac~=0.10.0 RUN aws s3 cp --no-sign-request s3://clay-model-ckpt/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index db5dbd21..ce4b69d9 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -70,10 +70,12 @@ def process_scene(clay, path, batchsize): "naip-analytic", str(path), fl, ExtraArgs={"RequestPayer": "requester"} ) - item = create_stac_item(fl.name, with_proj=True) - - item.datetime = date - item.id = f"{state}_{path.stem}" + item = create_stac_item( + fl.name, + with_proj=True, + input_datetime=date, + id=f"{state}_{path.stem}", + ) try: indexer = NoStatsChipIndexer(item) From 89f1be104ef3070525dc23d6932aa0a0aed5d6ac Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 5 Nov 2024 18:01:19 +0000 Subject: [PATCH 71/83] Fix datetime bug for files that have date stamps in them --- embeddings/Dockerfile | 7 +++---- embeddings/all-naip.py | 5 +++++ 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 26024c27..75555039 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -4,6 +4,9 @@ WORKDIR /model RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . +RUN aws s3 cp --no-sign-request s3://clay-model-ckpt/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt +RUN aws s3 cp --no-sign-request s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip + RUN pip install \ einops~=0.7.0 \ fiona~=1.9.5 \ @@ -32,10 +35,6 @@ RUN pip install \ wandb==0.17.5 \ rio_stac~=0.10.0 - -RUN aws s3 cp --no-sign-request s3://clay-model-ckpt/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt -RUN aws s3 cp --no-sign-request s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip - # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . RUN cp embeddings/all-sentinel.py . diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index ce4b69d9..09b94c07 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -70,11 +70,16 @@ def process_scene(clay, path, batchsize): "naip-analytic", str(path), fl, ExtraArgs={"RequestPayer": "requester"} ) + # Prepare properties, some NAIP imagery contains date stamps that + # raise an error in create_stac_item. + props = {"start_datetime": date, "end_datetime": date} + item = create_stac_item( fl.name, with_proj=True, input_datetime=date, id=f"{state}_{path.stem}", + properties=props, ) try: From eefae203ff5cca17881e6ace21077b63a8b3e033 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 6 Nov 2024 08:21:10 +0000 Subject: [PATCH 72/83] Add check for files that are already processed --- embeddings/all-naip.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 09b94c07..ceddc8d8 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -7,6 +7,7 @@ from pathlib import Path import boto3 +import botocore from rasterio.errors import RasterioIOError from rio_stac import create_stac_item from stacchip.chipper import Chipper @@ -117,6 +118,18 @@ def process_scene(clay, path, batchsize): write_to_table(embeddings=cls_embeddings, **kwargs) +def check_exists(path): + s3 = boto3.client("s3") + try: + s3.head_object( + Bucket=EMBEDDINGS_BUCKET, + Key=f"{path.parent}/{path.stem}.parquet", + ) + return True + except botocore.exceptions.ClientError: + return False + + def process(): if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") @@ -129,9 +142,14 @@ def process(): clay = load_clay() for i in range(index * items_per_job, (index + 1) * items_per_job): + scene = scenes[i] + if check_exists(scene): + logger.debug(f"Skipping scene because exists: {scene}") + continue + process_scene( clay=clay, - path=scenes[i], + path=scene, batchsize=batchsize, ) From 249695334427051227259c9a70072c14d741c189 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 6 Nov 2024 08:28:02 +0000 Subject: [PATCH 73/83] Move embeddings to CPU early --- embeddings/utils.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/embeddings/utils.py b/embeddings/utils.py index fba30304..7be73fd3 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -114,12 +114,14 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs unmsk_patch, unmsk_idx, msk_idx, msk_matrix = clay.model.encoder(datacube) # The first embedding is the class token, which is the # overall single embedding we want to keep. + batch_cls_embeddings = unmsk_patch[:, 0, :].cpu().numpy() + batch_patch_embeddings = unmsk_patch[:, 1:, :].cpu().numpy() if cls_embeddings is None: - cls_embeddings = unmsk_patch[:, 0, :] - patch_embeddings = unmsk_patch[:, 1:, :] + cls_embeddings = batch_cls_embeddings + patch_embeddings = batch_patch_embeddings else: - cls_embeddings = torch.vstack((cls_embeddings, unmsk_patch[:, 0, :])) - patch_embeddings = torch.vstack((patch_embeddings, unmsk_patch[:, 1:, :])) + cls_embeddings = np.vstack((cls_embeddings, batch_cls_embeddings)) + patch_embeddings = np.vstack((patch_embeddings, batch_patch_embeddings)) return cls_embeddings, patch_embeddings @@ -142,16 +144,15 @@ def load_clay(): def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): - np_embeddings = embeddings.cpu().numpy() index = {"geometry": ga.as_geoarrow([dat.wkt for dat in bboxs])} if len(embeddings.shape) == EMBEDDING_SHAPE_CLASS: # Handle class embeddings - index["embeddings"] = [np.ascontiguousarray(dat) for dat in np_embeddings] + index["embeddings"] = [np.ascontiguousarray(dat) for dat in embeddings] elif len(embeddings.shape) == EMBEDDING_SHAPE_PATCH: # Handle patch embeddings for i in range(embeddings.shape[1]): index[f"patch_embeddings_{i}"] = [ - np.ascontiguousarray(dat) for dat in np_embeddings[:, i, :] + np.ascontiguousarray(dat) for dat in embeddings[:, i, :] ] table = pa.table( From 097e291b16c89662609a7fd5704b9aa3c95d1f5c Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Wed, 6 Nov 2024 11:20:23 +0000 Subject: [PATCH 74/83] Remove patch embedding extraction --- embeddings/all-naip.py | 2 +- embeddings/all-sentinel.py | 3 +-- embeddings/utils.py | 5 +---- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index ceddc8d8..a7268220 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -97,7 +97,7 @@ def process_scene(clay, path, batchsize): mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=gsd ) # Embed data - cls_embeddings, patch_embeddings = get_embeddings( + cls_embeddings = get_embeddings( clay=clay, pixels_norm=pixels_norm, time_norm=time_norm, diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index f9b95bcb..94d4410e 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -59,7 +59,7 @@ def process_scene(clay, path, batchsize): mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=GSD ) # Embed data - cls_embeddings, patch_embeddings = get_embeddings( + cls_embeddings = get_embeddings( clay=clay, pixels_norm=pixels_norm, time_norm=time_norm, @@ -77,7 +77,6 @@ def process_scene(clay, path, batchsize): ) write_to_table(embeddings=cls_embeddings, **kwargs) - write_to_table(embeddings=patch_embeddings, **kwargs) def process(): diff --git a/embeddings/utils.py b/embeddings/utils.py index 7be73fd3..66435c2f 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -115,15 +115,12 @@ def get_embeddings(clay, pixels_norm, time_norm, latlon_norm, waves, gsd, batchs # The first embedding is the class token, which is the # overall single embedding we want to keep. batch_cls_embeddings = unmsk_patch[:, 0, :].cpu().numpy() - batch_patch_embeddings = unmsk_patch[:, 1:, :].cpu().numpy() if cls_embeddings is None: cls_embeddings = batch_cls_embeddings - patch_embeddings = batch_patch_embeddings else: cls_embeddings = np.vstack((cls_embeddings, batch_cls_embeddings)) - patch_embeddings = np.vstack((patch_embeddings, batch_patch_embeddings)) - return cls_embeddings, patch_embeddings + return cls_embeddings def load_clay(): From 0a6f476f40d7d6812e238770c1b612384b008406 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 12 Nov 2024 09:07:00 +0000 Subject: [PATCH 75/83] Sentinel-2 2024 run preparation --- embeddings/Dockerfile | 1 + embeddings/all-sentinel.py | 27 ++++++++++++++++++++++----- embeddings/utils.py | 18 +++++++++++++++++- 3 files changed, 40 insertions(+), 6 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 75555039..cf3b3d6f 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -6,6 +6,7 @@ RUN git clone -b all-of-naip https://github.com/Clay-foundation/model.git . RUN aws s3 cp --no-sign-request s3://clay-model-ckpt/v1.5.0-no-mrl-dinov2/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt RUN aws s3 cp --no-sign-request s3://clay-mgrs-samples/naip-manifest.txt.zip data/naip-manifest.txt.zip +RUN aws s3 cp --no-sign-request s3://clay-mgrs-samples/element84-tiles-2023.gz data/element84-tiles-2023.gz RUN pip install \ einops~=0.7.0 \ diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 94d4410e..84916ca8 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -18,8 +18,9 @@ write_to_table, ) +logging.basicConfig() logger = logging.getLogger("clay") - +logger.setLevel(logging.DEBUG) SCENES_LIST = "data/element84-tiles-2023.gz" EMBEDDINGS_BUCKET = "clay-embeddings-sentinel-2" @@ -29,11 +30,16 @@ def open_scenes_list(): with gzip.open(SCENES_LIST) as fl: data = fl.readlines() - return [dat.decode().rstrip() for dat in data] + data = [dat.decode().rstrip() for dat in data] + data = [dat for dat in data if dat.split("/")[7] == "2024"] + # Process the X, C, and D regions last + data = sorted(data, key=lambda dat: dat.split("/")[5] in ["X", "C", "D"]) + logger.debug(f"Found {len(data)} scenes to process") + return data def process_scene(clay, path, batchsize): - bands, waves, mean, std = load_metadata("sentinel_2_l2a") + bands, waves, mean, std = load_metadata("sentinel-2-l2a") key = path.replace("s3://sentinel-cogs/", "") @@ -43,11 +49,17 @@ def process_scene(clay, path, batchsize): item = Item.from_dict(stac_json) - bands, waves, mean, std = load_metadata("naip") + # Sanity checks + if "red" not in item.assets: + logger.debug(f"No red band for {key}") + return + elif not item.ext.has("proj"): + logger.debug(f"No proj for {key}") + return try: indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) - chipper = Chipper(item, assets=bands) + chipper = Chipper(indexer, assets=bands) bboxs, datetimes, pixels = get_pixels( item=item, indexer=indexer, chipper=chipper ) @@ -55,9 +67,14 @@ def process_scene(clay, path, batchsize): logger.warning("Skipping scene due to rasterio io error") return + if not len(pixels): + logger.debug("Finishing early, no valid data found in scene.") + return + time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=GSD ) + # Embed data cls_embeddings = get_embeddings( clay=clay, diff --git a/embeddings/utils.py b/embeddings/utils.py index 66435c2f..2208484d 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -17,6 +17,9 @@ EMBEDDING_SHAPE_CLASS = 2 EMBEDDING_SHAPE_PATCH = 3 +CLOUD_LIMIT = 0.1 +NODATA_LIMIT = 0.01 + logger = logging.getLogger("clay") @@ -76,7 +79,19 @@ def get_pixels(item, indexer, chipper): bboxs = [] chip_ids = [] item_ids = [] - for x, y, chip in chipper: + for index in range(len(chipper)): + y = index // chipper.indexer.x_size + x = index % chipper.indexer.x_size + + cloud_percentage, nodata_percentage = chipper.indexer.get_stats(x, y) + print(index, y, x, cloud_percentage, nodata_percentage) + if cloud_percentage > CLOUD_LIMIT: + continue + elif nodata_percentage > NODATA_LIMIT: + continue + + chip = chipper.chip(x, y) + chips.append(chip) datetimes.append(item.datetime) bboxs.append(indexer.get_chip_bbox(x, y)) @@ -84,6 +99,7 @@ def get_pixels(item, indexer, chipper): item_ids.append(item.id) pixels = np.array([np.array(list(chip.values())).squeeze() for chip in chips]) + return bboxs, datetimes, pixels From 6c867ee059f367d26f87ea55f05bdb378fea82c9 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 12 Nov 2024 11:03:47 +0000 Subject: [PATCH 76/83] Fix sentinel paths for output --- embeddings/Dockerfile | 2 +- embeddings/all-naip.py | 1 + embeddings/all-sentinel.py | 11 ++++++----- embeddings/utils.py | 6 ++++-- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index cf3b3d6f..7c591517 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -40,4 +40,4 @@ RUN pip install \ RUN cp embeddings/all-naip.py . RUN cp embeddings/all-sentinel.py . -ENTRYPOINT ["python", "all-naip.py"] +ENTRYPOINT ["python"] diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index a7268220..5a60b54f 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -113,6 +113,7 @@ def process_scene(clay, path, batchsize): gsd=gsd, destination_bucket=EMBEDDINGS_BUCKET, path=path, + source_bucket="naip-analytic", ) logger.debug("Writing class embeddings") write_to_table(embeddings=cls_embeddings, **kwargs) diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 84916ca8..fe495927 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -2,6 +2,7 @@ import json import logging import os +from pathlib import Path import boto3 from pystac import Item @@ -34,6 +35,7 @@ def open_scenes_list(): data = [dat for dat in data if dat.split("/")[7] == "2024"] # Process the X, C, and D regions last data = sorted(data, key=lambda dat: dat.split("/")[5] in ["X", "C", "D"]) + data = [Path(dat.replace("s3://sentinel-cogs/", "")) for dat in data] logger.debug(f"Found {len(data)} scenes to process") return data @@ -41,20 +43,18 @@ def open_scenes_list(): def process_scene(clay, path, batchsize): bands, waves, mean, std = load_metadata("sentinel-2-l2a") - key = path.replace("s3://sentinel-cogs/", "") - s3 = boto3.resource("s3") - stac_json = json.load(s3.Object("sentinel-cogs", key).get()["Body"]) + stac_json = json.load(s3.Object("sentinel-cogs", str(path)).get()["Body"]) item = Item.from_dict(stac_json) # Sanity checks if "red" not in item.assets: - logger.debug(f"No red band for {key}") + logger.debug(f"No red band for {path}") return elif not item.ext.has("proj"): - logger.debug(f"No proj for {key}") + logger.debug(f"No proj for {path}") return try: @@ -91,6 +91,7 @@ def process_scene(clay, path, batchsize): gsd=gsd, destination_bucket=EMBEDDINGS_BUCKET, path=path, + source_bucket="sentinel-cogs", ) write_to_table(embeddings=cls_embeddings, **kwargs) diff --git a/embeddings/utils.py b/embeddings/utils.py index 2208484d..05b0ca8d 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -156,7 +156,9 @@ def load_clay(): return model.to(device) -def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): +def write_to_table( # noqa: PLR0913 + embeddings, bboxs, datestr, gsd, destination_bucket, path, source_bucket +): index = {"geometry": ga.as_geoarrow([dat.wkt for dat in bboxs])} if len(embeddings.shape) == EMBEDDING_SHAPE_CLASS: # Handle class embeddings @@ -173,7 +175,7 @@ def write_to_table(embeddings, bboxs, datestr, gsd, destination_bucket, path): metadata={ "date": datestr, "gsd": str(gsd[0]), - "uri": f"s3://naip-analytic/{path}", + "uri": f"s3://{source_bucket}/{path}", }, ) From ac1434be70fb2b342a1b96da83749460a878d904 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Fri, 22 Nov 2024 10:36:27 +0000 Subject: [PATCH 77/83] Pre-download S2 scene, batch pixel load and embeding generation --- embeddings/all-sentinel.py | 95 ++++++++++++++++++++++++++------------ embeddings/utils.py | 8 +++- 2 files changed, 71 insertions(+), 32 deletions(-) diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index fe495927..875188b6 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -2,18 +2,18 @@ import json import logging import os +import tempfile from pathlib import Path import boto3 +import numpy as np from pystac import Item -from rasterio.errors import RasterioIOError from stacchip.chipper import Chipper from stacchip.indexer import Sentinel2Indexer from embeddings.utils import ( get_embeddings, get_pixels, - load_clay, load_metadata, prepare_datacube, write_to_table, @@ -26,6 +26,7 @@ SCENES_LIST = "data/element84-tiles-2023.gz" EMBEDDINGS_BUCKET = "clay-embeddings-sentinel-2" GSD = 10 +S2_BUCKET = "sentinel-2-cogs" def open_scenes_list(): @@ -40,6 +41,21 @@ def open_scenes_list(): return data +def download_scenes_local(tmp, item, bands): + s3 = boto3.client("s3") + for band in bands: + local_asset_path = f"{tmp}/{band}.tif" + remote_asset_key = item.assets[band].href.replace( + "https://sentinel-cogs.s3.us-west-2.amazonaws.com/", "" + ) + print(f"Downloading band {band} to {local_asset_path}") + with open(local_asset_path, mode="w+b") as fl: + s3.download_fileobj("sentinel-cogs", remote_asset_key, fl) + item.assets[band].href = local_asset_path + + return item + + def process_scene(clay, path, batchsize): bands, waves, mean, std = load_metadata("sentinel-2-l2a") @@ -57,36 +73,54 @@ def process_scene(clay, path, batchsize): logger.debug(f"No proj for {path}") return - try: + all_bboxs = [] + all_cls_embeddings = None + + with tempfile.TemporaryDirectory() as tmp: + item = download_scenes_local(tmp, item, bands) indexer = Sentinel2Indexer(item, chip_max_nodata=0.1) chipper = Chipper(indexer, assets=bands) - bboxs, datetimes, pixels = get_pixels( - item=item, indexer=indexer, chipper=chipper - ) - except RasterioIOError: - logger.warning("Skipping scene due to rasterio io error") - return + logger.debug(f"Creating chips for {item.id}") + STEP = 50 + for index in range(0, len(chipper), STEP): + bboxs, datetimes, pixels = get_pixels( + item=item, + indexer=indexer, + chipper=chipper, + start=index, + end=index + STEP, + ) + + if not len(pixels): + continue + + time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( + mean=mean, + std=std, + datetimes=datetimes, + bboxs=bboxs, + pixels=pixels, + gsd=GSD, + ) + + # Embed data + cls_embeddings = get_embeddings( + clay=clay, + pixels_norm=pixels_norm, + time_norm=time_norm, + latlon_norm=latlon_norm, + waves=waves, + gsd=gsd, + batchsize=batchsize, + ) + all_bboxs += bboxs + if all_cls_embeddings is None: + all_cls_embeddings = cls_embeddings + else: + all_cls_embeddings = np.vstack((all_cls_embeddings, cls_embeddings)) - if not len(pixels): - logger.debug("Finishing early, no valid data found in scene.") - return - - time_norm, latlon_norm, gsd, pixels_norm = prepare_datacube( - mean=mean, std=std, datetimes=datetimes, bboxs=bboxs, pixels=pixels, gsd=GSD - ) - - # Embed data - cls_embeddings = get_embeddings( - clay=clay, - pixels_norm=pixels_norm, - time_norm=time_norm, - latlon_norm=latlon_norm, - waves=waves, - gsd=gsd, - batchsize=batchsize, - ) kwargs = dict( - bboxs=bboxs, + bboxs=all_bboxs, datestr=str(item.datetime.date()), gsd=gsd, destination_bucket=EMBEDDINGS_BUCKET, @@ -94,7 +128,7 @@ def process_scene(clay, path, batchsize): source_bucket="sentinel-cogs", ) - write_to_table(embeddings=cls_embeddings, **kwargs) + write_to_table(embeddings=all_cls_embeddings, **kwargs) def process(): @@ -105,7 +139,8 @@ def process(): batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) scenes = open_scenes_list() - clay = load_clay() + # clay = load_clay() + clay = None for i in range(index * items_per_job, (index + 1) * items_per_job): process_scene( diff --git a/embeddings/utils.py b/embeddings/utils.py index 05b0ca8d..45943d99 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -73,13 +73,17 @@ def prepare_datacube(mean, std, datetimes, bboxs, pixels, gsd): return time_norm, latlon_norm, gsd, pixels_norm -def get_pixels(item, indexer, chipper): +def get_pixels(item, indexer, chipper, start=None, end=None): chips = [] datetimes = [] bboxs = [] chip_ids = [] item_ids = [] - for index in range(len(chipper)): + if start: + index_range = range(start, min(end, len(chipper))) + else: + index_range = range(len(chipper)) + for index in index_range: y = index // chipper.indexer.x_size x = index % chipper.indexer.x_size From c270a4c2dc388fdd6829b9987d57dac07277efa3 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Fri, 22 Nov 2024 15:16:35 +0000 Subject: [PATCH 78/83] Use torch tensor for normalization --- embeddings/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/utils.py b/embeddings/utils.py index 45943d99..00b977a3 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -68,7 +68,7 @@ def prepare_datacube(mean, std, datetimes, bboxs, pixels, gsd): gsd = [gsd] - pixels_norm = transform(pixels) + pixels_norm = transform(torch.tensor(pixels, dtype=torch.float32)).numpy() return time_norm, latlon_norm, gsd, pixels_norm From 88ed3c0cee44bf323e67e06eea75838a967a33bd Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Fri, 22 Nov 2024 16:43:04 +0000 Subject: [PATCH 79/83] Use custom endpoint on demand --- embeddings/all-naip.py | 14 +++++++++++--- embeddings/utils.py | 12 +++++++++++- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 5a60b54f..57fc3af0 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -28,7 +28,7 @@ MANIFEST = "data/naip-manifest.txt.zip" -EMBEDDINGS_BUCKET = "clay-embeddings-naip" +EMBEDDINGS_BUCKET = os.environ["EMBEDDINGS_BUCKET"] def open_scene_list(limit_to_state=None): @@ -120,7 +120,15 @@ def process_scene(clay, path, batchsize): def check_exists(path): - s3 = boto3.client("s3") + if "ENDPOINT_URL" in os.environ: + s3 = boto3.client( + "s3", + endpoint_url=os.environ.get("ENDPOINT_URL"), + aws_access_key_id=os.environ.get("ENDPOINT_KEY_ID"), + aws_secret_access_key=os.environ.get("ENDPOINT_ACCESS_KEY"), + ) + else: + s3 = boto3.client("s3") try: s3.head_object( Bucket=EMBEDDINGS_BUCKET, @@ -135,7 +143,7 @@ def process(): if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") index = int(os.environ.get("AWS_BATCH_JOB_ARRAY_INDEX", 0)) - items_per_job = int(os.environ.get("ITEMS_PER_JOB", 2)) + items_per_job = int(os.environ.get("ITEMS_PER_JOB", 100)) batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) limit_to_state = os.environ.get("LIMIT_TO_STATE", None) diff --git a/embeddings/utils.py b/embeddings/utils.py index 00b977a3..37e03a36 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -1,5 +1,6 @@ import logging import math +import os import boto3 import geoarrow.pyarrow as ga @@ -186,7 +187,16 @@ def write_to_table( # noqa: PLR0913 writer = pa.BufferOutputStream() gaio.write_geoparquet_table(table, writer) body = bytes(writer.getvalue()) - s3_resource = boto3.resource("s3") + if "ENDPOINT_URL" in os.environ: + s3_resource = boto3.resource( + "s3", + endpoint_url=os.environ.get("ENDPOINT_URL"), + aws_access_key_id=os.environ.get("ENDPOINT_KEY_ID"), + aws_secret_access_key=os.environ.get("ENDPOINT_ACCESS_KEY"), + ) + else: + s3_resource = boto3.resource("s3") + s3_bucket = s3_resource.Bucket(name=destination_bucket) s3_bucket.put_object( Body=body, From ceecb6138705cb28a5f4d3f61f22b19a2f625edb Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Sun, 24 Nov 2024 14:47:05 +0000 Subject: [PATCH 80/83] Check exists for Sentinel-2 process --- embeddings/README.md | 8 ++++++++ embeddings/all-naip.py | 31 ++++++++----------------------- embeddings/all-sentinel.py | 9 +++++++-- embeddings/utils.py | 24 +++++++++++++++++++++++- 4 files changed, 46 insertions(+), 26 deletions(-) diff --git a/embeddings/README.md b/embeddings/README.md index cc7685ee..f8a835f3 100644 --- a/embeddings/README.md +++ b/embeddings/README.md @@ -11,6 +11,14 @@ to choose which files from the archives to process. This is set automatically by AWS Batch when using array jobs. Outside of array jobs, this index variable needs to be specified manually. +The script also requires the `EMBEDDINGS_BUCKET` environment variable, +specifying the name of the output bucket. + +To specify a custom bucket location (for source coop for instance), use the +`ENDPOINT_URL`, `ENDPOINT_KEY_ID`, and `ENDPOINT_ACCESS_KEY` environment +variables. + + ### Build docker image Embedding runs are dockerized for parallel computing. To build the docker image diff --git a/embeddings/all-naip.py b/embeddings/all-naip.py index 57fc3af0..69594251 100644 --- a/embeddings/all-naip.py +++ b/embeddings/all-naip.py @@ -7,13 +7,13 @@ from pathlib import Path import boto3 -import botocore from rasterio.errors import RasterioIOError from rio_stac import create_stac_item from stacchip.chipper import Chipper from stacchip.indexer import NoStatsChipIndexer from embeddings.utils import ( + check_exists, get_embeddings, get_pixels, load_clay, @@ -29,6 +29,7 @@ MANIFEST = "data/naip-manifest.txt.zip" EMBEDDINGS_BUCKET = os.environ["EMBEDDINGS_BUCKET"] +HOUR_OF_DAY = 12 def open_scene_list(limit_to_state=None): @@ -59,7 +60,9 @@ def process_scene(clay, path, batchsize): """ state = path.parts[0] datestr = path.stem.split("_")[-1] - date = datetime.datetime(int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8])) + date = datetime.datetime( + int(datestr[:4]), int(datestr[4:6]), int(datestr[6:8]), HOUR_OF_DAY + ) gsd = float(path.parts[2].replace("cm", "")) / 100 bands, waves, mean, std = load_metadata("naip") @@ -87,7 +90,9 @@ def process_scene(clay, path, batchsize): indexer = NoStatsChipIndexer(item) chipper = Chipper(indexer) bboxs, datetimes, pixels = get_pixels( - item=item, indexer=indexer, chipper=chipper + item=item, + indexer=indexer, + chipper=chipper, ) except RasterioIOError: logger.warning("Skipping scene due to rasterio io error") @@ -119,26 +124,6 @@ def process_scene(clay, path, batchsize): write_to_table(embeddings=cls_embeddings, **kwargs) -def check_exists(path): - if "ENDPOINT_URL" in os.environ: - s3 = boto3.client( - "s3", - endpoint_url=os.environ.get("ENDPOINT_URL"), - aws_access_key_id=os.environ.get("ENDPOINT_KEY_ID"), - aws_secret_access_key=os.environ.get("ENDPOINT_ACCESS_KEY"), - ) - else: - s3 = boto3.client("s3") - try: - s3.head_object( - Bucket=EMBEDDINGS_BUCKET, - Key=f"{path.parent}/{path.stem}.parquet", - ) - return True - except botocore.exceptions.ClientError: - return False - - def process(): if "AWS_BATCH_JOB_ARRAY_INDEX" not in os.environ: raise ValueError("AWS_BATCH_JOB_ARRAY_INDEX env var not set") diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 875188b6..4015e090 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -12,8 +12,10 @@ from stacchip.indexer import Sentinel2Indexer from embeddings.utils import ( + check_exists, get_embeddings, get_pixels, + load_clay, load_metadata, prepare_datacube, write_to_table, @@ -139,10 +141,13 @@ def process(): batchsize = int(os.environ.get("EMBEDDING_BATCH_SIZE", 50)) scenes = open_scenes_list() - # clay = load_clay() - clay = None + clay = load_clay() for i in range(index * items_per_job, (index + 1) * items_per_job): + if check_exists(scenes[i]): + logger.debug(f"Skipping scene because exists: {scenes[i]}") + continue + process_scene( clay=clay, path=scenes[i], diff --git a/embeddings/utils.py b/embeddings/utils.py index 37e03a36..8298293b 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -3,6 +3,7 @@ import os import boto3 +import botocore import geoarrow.pyarrow as ga import numpy as np import pyarrow as pa @@ -17,6 +18,7 @@ CHECKPOINT = "data/mae_v1.5.0_epoch-07_val-loss-0.1718.ckpt" EMBEDDING_SHAPE_CLASS = 2 EMBEDDING_SHAPE_PATCH = 3 +EMBEDDINGS_BUCKET = os.environ["EMBEDDINGS_BUCKET"] CLOUD_LIMIT = 0.1 NODATA_LIMIT = 0.01 @@ -24,6 +26,26 @@ logger = logging.getLogger("clay") +def check_exists(path): + if "ENDPOINT_URL" in os.environ: + s3 = boto3.client( + "s3", + endpoint_url=os.environ.get("ENDPOINT_URL"), + aws_access_key_id=os.environ.get("ENDPOINT_KEY_ID"), + aws_secret_access_key=os.environ.get("ENDPOINT_ACCESS_KEY"), + ) + else: + s3 = boto3.client("s3") + try: + s3.head_object( + Bucket=EMBEDDINGS_BUCKET, + Key=f"{path.parent}/{path.stem}.parquet", + ) + return True + except botocore.exceptions.ClientError: + return False + + def load_metadata(platform): metadata = Box(yaml.safe_load(open("configs/metadata.yaml"))) platform_meta = getattr(metadata, platform) @@ -89,7 +111,7 @@ def get_pixels(item, indexer, chipper, start=None, end=None): x = index % chipper.indexer.x_size cloud_percentage, nodata_percentage = chipper.indexer.get_stats(x, y) - print(index, y, x, cloud_percentage, nodata_percentage) + if cloud_percentage > CLOUD_LIMIT: continue elif nodata_percentage > NODATA_LIMIT: From 55a82a8308b14e9933b8628380c656ff385e5859 Mon Sep 17 00:00:00 2001 From: Daniel Wiesmann Date: Tue, 26 Nov 2024 14:54:38 +0000 Subject: [PATCH 81/83] Make Dockerfile cacheable --- embeddings/Dockerfile | 2 ++ embeddings/all-sentinel.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/embeddings/Dockerfile b/embeddings/Dockerfile index 7c591517..eb532ab7 100644 --- a/embeddings/Dockerfile +++ b/embeddings/Dockerfile @@ -36,6 +36,8 @@ RUN pip install \ wandb==0.17.5 \ rio_stac~=0.10.0 +RUN git pull && git checkout ceecb6138705cb28a5f4d3f61f22b19a2f625edb + # Move file to home directory so that relative imports work RUN cp embeddings/all-naip.py . RUN cp embeddings/all-sentinel.py . diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 4015e090..2e5b70d9 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -124,7 +124,7 @@ def process_scene(clay, path, batchsize): kwargs = dict( bboxs=all_bboxs, datestr=str(item.datetime.date()), - gsd=gsd, + gsd=GSD, destination_bucket=EMBEDDINGS_BUCKET, path=path, source_bucket="sentinel-cogs", From 085f3718df45f821e2f0cdfe8d27f62fa033aba6 Mon Sep 17 00:00:00 2001 From: Mason Grimshaw Date: Fri, 20 Dec 2024 21:17:35 -0700 Subject: [PATCH 82/83] Update all-sentinel.py Quick change for Brazil run --- embeddings/all-sentinel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/embeddings/all-sentinel.py b/embeddings/all-sentinel.py index 2e5b70d9..505d273e 100644 --- a/embeddings/all-sentinel.py +++ b/embeddings/all-sentinel.py @@ -25,7 +25,7 @@ logger = logging.getLogger("clay") logger.setLevel(logging.DEBUG) -SCENES_LIST = "data/element84-tiles-2023.gz" +SCENES_LIST = "data/element84-tiles-2023-brazil.gz" EMBEDDINGS_BUCKET = "clay-embeddings-sentinel-2" GSD = 10 S2_BUCKET = "sentinel-2-cogs" @@ -112,7 +112,7 @@ def process_scene(clay, path, batchsize): time_norm=time_norm, latlon_norm=latlon_norm, waves=waves, - gsd=gsd, + gsd=GSD, batchsize=batchsize, ) all_bboxs += bboxs From 2a48007c29054f793234dcd863d48d3bda1afabc Mon Sep 17 00:00:00 2001 From: Mason Grimshaw Date: Sat, 21 Dec 2024 00:16:25 -0700 Subject: [PATCH 83/83] Update utils.py Maybe a bug when writing metadata? --- embeddings/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embeddings/utils.py b/embeddings/utils.py index 8298293b..61cc3919 100644 --- a/embeddings/utils.py +++ b/embeddings/utils.py @@ -201,7 +201,7 @@ def write_to_table( # noqa: PLR0913 index, metadata={ "date": datestr, - "gsd": str(gsd[0]), + "gsd": str(gsd), "uri": f"s3://{source_bucket}/{path}", }, )