rslearn 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

rslearn/lightning_cli.py

@@ -1,12 +1,107 @@
 """LightningCLI for rslearn."""
 
+import hashlib
+import json
+import os
+import shutil
 import sys
+import tempfile
 
+import fsspec
+import jsonargparse
+import wandb
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
 from lightning.pytorch.cli import LightningArgumentParser, LightningCLI
+from lightning.pytorch.utilities import rank_zero_only
+from upath import UPath
 
 from rslearn.arg_parser import RslearnArgumentParser
+from rslearn.log_utils import get_logger
 from rslearn.train.data_module import RslearnDataModule
 from rslearn.train.lightning_module import RslearnLightningModule
+from rslearn.utils.fsspec import open_atomic
+
+WANDB_ID_FNAME = "wandb_id"
+
+logger = get_logger(__name__)
+
+
+def get_cached_checkpoint(checkpoint_fname: UPath) -> str:
+    """Get a local cached version of the specified checkpoint.
+
+    If checkpoint_fname is already local, then it is returned. Otherwise, it is saved
+    in a deterministic local cache directory under the system temporary directory, and
+    the cached filename is returned.
+
+    Note that the cache is not deleted when the program exits.
+
+    Args:
+        checkpoint_fname: the potentially non-local checkpoint file to load.
+
+    Returns:
+        a local filename containing the same checkpoint.
+    """
+    is_local = isinstance(
+        checkpoint_fname.fs, fsspec.implementations.local.LocalFileSystem
+    )
+    if is_local:
+        return checkpoint_fname.path
+
+    cache_id = hashlib.sha256(str(checkpoint_fname).encode()).hexdigest()
+    local_fname = os.path.join(
+        tempfile.gettempdir(), "rslearn_cache", "checkpoints", f"{cache_id}.ckpt"
+    )
+
+    if os.path.exists(local_fname):
+        logger.info(
+            "using cached checkpoint for %s at %s", str(checkpoint_fname), local_fname
+        )
+        return local_fname
+
+    logger.info("caching checkpoint %s to %s", str(checkpoint_fname), local_fname)
+    os.makedirs(os.path.dirname(local_fname), exist_ok=True)
+    with checkpoint_fname.open("rb") as src:
+        with open_atomic(UPath(local_fname), "wb") as dst:
+            shutil.copyfileobj(src, dst)
+
+    return local_fname
+
+
+class SaveWandbRunIdCallback(Callback):
+    """Callback to save the wandb run ID to project directory in case of resume."""
+
+    def __init__(
+        self,
+        project_dir: str,
+        config_str: str,
+    ) -> None:
+        """Create a new SaveWandbRunIdCallback.
+
+        Args:
+            project_dir: the project directory.
+            config_str: the JSON-encoded configuration of this experiment
+        """
+        self.project_dir = project_dir
+        self.config_str = config_str
+
+    @rank_zero_only
+    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        """Called just before fit starts I think.
+
+        Args:
+            trainer: the Trainer object.
+            pl_module: the LightningModule object.
+        """
+        wandb_id = wandb.run.id
+
+        project_dir = UPath(self.project_dir)
+        project_dir.mkdir(parents=True, exist_ok=True)
+        with (project_dir / WANDB_ID_FNAME).open("w") as f:
+            f.write(wandb_id)
+
+        if self.config_str is not None and "project_name" not in wandb.config:
+            wandb.config.update(json.loads(self.config_str))
 
 
 class RslearnLightningCLI(LightningCLI):
@@ -23,6 +118,266 @@ class RslearnLightningCLI(LightningCLI):
             "data.init_args.task", "model.init_args.task", apply_on="instantiate"
         )
 
+        # Project management option to have rslearn manage checkpoints and W&B run.
+        parser.add_argument(
+            "--management_dir",
+            type=str | None,
+            help="Enable project management, and use this directory to store checkpoints and configs. If enabled, rslearn will automatically manages checkpoint directory/loading and W&B run",
+            default=None,
+        )
+        parser.add_argument(
+            "--project_name",
+            type=str | None,
+            help="The project name (used with --management_dir)",
+            default=None,
+        )
+        parser.add_argument(
+            "--run_name",
+            type=str | None,
+            help="A unique name for this experiment (used with --management_dir)",
+            default=None,
+        )
+        parser.add_argument(
+            "--run_description",
+            type=str,
+            help="Optional description of this experiment (used with --management_dir)",
+            default="",
+        )
+        parser.add_argument(
+            "--load_checkpoint_mode",
+            type=str,
+            help="Which checkpoint to load, if any (used with --management_dir). 'none' never loads any checkpoint, 'last' loads the most recent checkpoint, and 'best' loads the best checkpoint. 'auto' will use 'last' during fit and 'best' during val/test/predict.",
+            default="auto",
+        )
+        parser.add_argument(
+            "--load_checkpoint_required",
+            type=str,
+            help="Whether to fail if the expected checkpoint based on load_checkpoint_mode does not exist (used with --management_dir). 'yes' will fail while 'no' won't. 'auto' will use 'no' during fit and 'yes' during val/test/predict.",
+            default="auto",
+        )
+        parser.add_argument(
+            "--log_mode",
+            type=str,
+            help="Whether to log to W&B (used with --management_dir). 'yes' will enable logging, 'no' will disable logging, and 'auto' will use 'yes' during fit and 'no' during val/test/predict.",
+            default="auto",
+        )
+
+    def _get_checkpoint_path(
+        self,
+        project_dir: UPath,
+        load_checkpoint_mode: str,
+        load_checkpoint_required: str,
+        stage: str,
+    ) -> str | None:
+        """Get path to checkpoint to load from, or None to not restore checkpoint.
+
+        Args:
+            project_dir: the project directory determined from the project management
+                directory.
+            load_checkpoint_mode: "none" to not load any checkpoint, "last" to load the
+                most recent checkpoint, "best" to load the best checkpoint. "auto" to
+                use "last" during fit and "best" during val/test/predict.
+            load_checkpoint_required: "yes" to fail if no checkpoint exists, "no" to
+                ignore. "auto" will use "no" during fit and "yes" during
+                val/test/predict.
+            stage: the lightning stage (fit/val/test/predict).
+
+        Returns:
+            the path to the checkpoint for setting c.ckpt_path, or None if no
+                checkpoint should be restored.
+        """
+        # Resolve auto options if used.
+        if load_checkpoint_mode == "auto":
+            if stage == "fit":
+                load_checkpoint_mode = "last"
+            else:
+                load_checkpoint_mode = "best"
+        if load_checkpoint_required == "auto":
+            if stage == "fit":
+                load_checkpoint_required = "no"
+            else:
+                load_checkpoint_required = "yes"
+
+        if load_checkpoint_required == "yes" and load_checkpoint_mode == "none":
+            raise ValueError(
+                "load_checkpoint_required cannot be set when load_checkpoint_mode is none"
+            )
+
+        ckpt_path: str | None = None
+
+        if load_checkpoint_mode == "best":
+            # Checkpoints should be either:
+            # - last.ckpt
+            # - of the form "A=B-C=D-....ckpt" with one key being epoch=X
+            # So we want the one with the highest epoch, and only use last.ckpt if
+            # it's the only option.
+            # User should set save_top_k=1 so there's just one, otherwise we won't
+            # actually know which one is the best.
+            best_checkpoint = None
+            best_epochs = None
+            for option in project_dir.iterdir():
+                if not option.name.endswith(".ckpt"):
+                    continue
+
+                # Try to see what epochs this checkpoint is at.
+                # If it is some other format, then set it 0 so we only use it if it's
+                # the only option.
+                # If it is last.ckpt then we set it -100 to only use it if there is not
+                # even another format like "best.ckpt".
+                extracted_epochs = 0
+                if option.name == "last.ckpt":
+                    extracted_epochs = -100
+
+                parts = option.name.split(".ckpt")[0].split("-")
+                for part in parts:
+                    kv_parts = part.split("=")
+                    if len(kv_parts) != 2:
+                        continue
+                    if kv_parts[0] != "epoch":
+                        continue
+                    extracted_epochs = int(kv_parts[1])
+
+                if best_epochs is None or extracted_epochs > best_epochs:
+                    best_checkpoint = option
+                    best_epochs = extracted_epochs
+
+            if best_checkpoint is not None:
+                # Cache the checkpoint so we only need to download once in case we
+                # reuse it later.
+                # We only cache with --load_best since this is the only scenario where we
+                # expect to keep reusing the same checkpoint.
+                ckpt_path = get_cached_checkpoint(best_checkpoint)
+
+        elif load_checkpoint_mode == "last":
+            last_checkpoint_path = project_dir / "last.ckpt"
+            if last_checkpoint_path.exists():
+                ckpt_path = str(last_checkpoint_path)
+
+        else:
+            raise ValueError(f"unknown load_checkpoint_mode {load_checkpoint_mode}")
+
+        if load_checkpoint_required == "yes" and ckpt_path is None:
+            raise ValueError(
+                "load_checkpoint_required is set but no checkpoint was found"
+            )
+
+        return ckpt_path
+
+    def enable_project_management(self, management_dir: str) -> None:
+        """Enable project management in the specified directory.
+
+        Args:
+            management_dir: the directory to store checkpoints and W&B.
+        """
+        subcommand = self.config.subcommand
+        c = self.config[subcommand]
+
+        # Project name and run name are required with project management.
+        if not c.project_name or not c.run_name:
+            raise ValueError(
+                "project name and run name must be set when using project management"
+            )
+
+        # Get project directory within the project management directory.
+        project_dir = UPath(management_dir) / c.project_name / c.run_name
+
+        # Add the W&B logger if it isn't already set, and (re-)configure it.
+        should_log = False
+        if c.log_mode == "yes":
+            should_log = True
+        elif c.log_mode == "auto":
+            should_log = subcommand == "fit"
+        if should_log:
+            if not c.trainer.logger:
+                c.trainer.logger = jsonargparse.Namespace(
+                    {
+                        "class_path": "lightning.pytorch.loggers.WandbLogger",
+                        "init_args": jsonargparse.Namespace(),
+                    }
+                )
+            c.trainer.logger.init_args.project = c.project_name
+            c.trainer.logger.init_args.name = c.run_name
+            if c.run_description:
+                c.trainer.logger.init_args.notes = c.run_description
+
+            # Add callback to save config to W&B.
+            upload_wandb_callback = None
+            if "callbacks" in c.trainer and c.trainer.callbacks:
+                for existing_callback in c.trainer.callbacks:
+                    if existing_callback.class_path == "SaveWandbRunIdCallback":
+                        upload_wandb_callback = existing_callback
+            else:
+                c.trainer.callbacks = []
+
+            if not upload_wandb_callback:
+                config_str = json.dumps(
+                    c.as_dict(), default=lambda _: "<not serializable>"
+                )
+                upload_wandb_callback = jsonargparse.Namespace(
+                    {
+                        "class_path": "SaveWandbRunIdCallback",
+                        "init_args": jsonargparse.Namespace(
+                            {
+                                "project_dir": str(project_dir),
+                                "config_str": config_str,
+                            }
+                        ),
+                    }
+                )
+                c.trainer.callbacks.append(upload_wandb_callback)
+        else:
+            c.trainer.logger = jsonargparse.Namespace({})
+
+        if subcommand == "fit":
+            # Set the checkpoint directory to match the project directory.
+            checkpoint_callback = None
+            if "callbacks" in c.trainer and c.trainer.callbacks:
+                for existing_callback in c.trainer.callbacks:
+                    if (
+                        existing_callback.class_path
+                        == "lightning.pytorch.callbacks.ModelCheckpoint"
+                    ):
+                        checkpoint_callback = existing_callback
+            else:
+                c.trainer.callbacks = []
+
+            if not checkpoint_callback:
+                checkpoint_callback = jsonargparse.Namespace(
+                    {
+                        "class_path": "lightning.pytorch.callbacks.ModelCheckpoint",
+                        "init_args": jsonargparse.Namespace(
+                            {
+                                "save_last": True,
+                                "save_top_k": 1,
+                                "monitor": "val_loss",
+                            }
+                        ),
+                    }
+                )
+                c.trainer.callbacks.append(checkpoint_callback)
+            checkpoint_callback.init_args.dirpath = str(project_dir)
+
+        # Load existing checkpoint.
+        checkpoint_path = self._get_checkpoint_path(
+            project_dir=project_dir,
+            load_checkpoint_mode=c.load_checkpoint_mode,
+            load_checkpoint_required=c.load_checkpoint_required,
+            stage=subcommand,
+        )
+        if checkpoint_path is not None:
+            logger.info(f"found checkpoint to resume from at {checkpoint_path}")
+            c.ckpt_path = checkpoint_path
+
+            # If we are resuming from a checkpoint for training, we also try to resume the W&B run.
+            if (
+                subcommand == "fit"
+                and (project_dir / WANDB_ID_FNAME).exists()
+                and should_log
+            ):
+                with (project_dir / WANDB_ID_FNAME).open("r") as f:
+                    wandb_id = f.read().strip()
+                    c.trainer.logger.init_args.id = wandb_id
+
     def before_instantiate_classes(self) -> None:
         """Called before Lightning class initialization.
 
@@ -33,7 +388,7 @@ class RslearnLightningCLI(LightningCLI):
 
         # If there is a RslearnPredictionWriter, set its path.
         prediction_writer_callback = None
-        if "callbacks" in c.trainer:
+        if "callbacks" in c.trainer and c.trainer.callbacks:
             for existing_callback in c.trainer.callbacks:
                 if (
                     existing_callback.class_path
@@ -53,6 +408,20 @@ class RslearnLightningCLI(LightningCLI):
         if subcommand == "predict":
             c.return_predictions = False
 
+        # For now we use DDP strategy with find_unused_parameters=True.
+        if subcommand == "fit":
+            c.trainer.strategy = jsonargparse.Namespace(
+                {
+                    "class_path": "lightning.pytorch.strategies.DDPStrategy",
+                    "init_args": jsonargparse.Namespace(
+                        {"find_unused_parameters": True}
+                    ),
+                }
+            )
+
+        if c.management_dir:
+            self.enable_project_management(c.management_dir)
+
 
 def model_handler() -> None:
     """Handler for any rslearn model X commands."""

rslearn/main.py

@@ -15,7 +15,7 @@ from upath import UPath
 
 from rslearn.config import LayerConfig
 from rslearn.const import WGS84_EPSG
-from rslearn.data_sources import Item, data_source_from_config
+from rslearn.data_sources import Item
 from rslearn.dataset import Dataset, Window, WindowLayerData
 from rslearn.dataset.add_windows import add_windows_from_box, add_windows_from_file
 from rslearn.dataset.handler_summaries import (
@@ -544,7 +544,7 @@ class IngestHandler:
                 tile_store, layer_name, layer_cfg
             )
             layer_cfg = self.dataset.layers[layer_name]
-            data_source = data_source_from_config(layer_cfg, self.dataset.path)
+            data_source = layer_cfg.instantiate_data_source(self.dataset.path)
 
             attempts_counter = AttemptsCounter()
             ingest_counts: IngestCounts | UnknownIngestCounts
@@ -640,7 +640,7 @@ class IngestHandler:
             if not layer_cfg.data_source.ingest:
                 continue
 
-            data_source = data_source_from_config(layer_cfg, self.dataset.path)
+            data_source = layer_cfg.instantiate_data_source(self.dataset.path)
 
             geometries_by_item: dict = {}
             for window, layer_datas in windows_and_layer_datas:

rslearn/models/clay/clay.py

@@ -8,6 +8,7 @@ from importlib.resources import files
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 import yaml
 from einops import rearrange
 from huggingface_hub import hf_hub_download
@@ -30,6 +31,7 @@ PATCH_SIZE = 8
 CLAY_MODALITIES = ["sentinel-2-l2a", "sentinel-1-rtc", "landsat-c2l1", "naip"]
 CONFIG_DIR = files("rslearn.models.clay.configs")
 CLAY_METADATA_PATH = str(CONFIG_DIR / "metadata.yaml")
+DEFAULT_IMAGE_RESOLUTION = 128  # image resolution during pretraining
 
 
 def get_clay_checkpoint_path(
@@ -49,6 +51,7 @@ class Clay(torch.nn.Module):
         modality: str = "sentinel-2-l2a",
         checkpoint_path: str | None = None,
         metadata_path: str = CLAY_METADATA_PATH,
+        do_resizing: bool = False,
     ) -> None:
         """Initialize the Clay model.
 
@@ -57,6 +60,7 @@ class Clay(torch.nn.Module):
             modality: The modality to use (subset of CLAY_MODALITIES).
             checkpoint_path: Path to clay-v1.5.ckpt, if None, fetch from HF Hub.
             metadata_path: Path to metadata.yaml.
+            do_resizing: Whether to resize the image to the input resolution.
         """
         super().__init__()
 
@@ -95,6 +99,14 @@ class Clay(torch.nn.Module):
 
         self.model_size = model_size
         self.modality = modality
+        self.do_resizing = do_resizing
+
+    def _resize_image(self, image: torch.Tensor, original_hw: int) -> torch.Tensor:
+        """Resize the image to the input resolution."""
+        new_hw = self.patch_size if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
+        return F.interpolate(
+            image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
+        )
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Forward pass for the Clay model.
@@ -114,7 +126,8 @@ class Clay(torch.nn.Module):
         chips = torch.stack(
             [inp[self.modality] for inp in inputs], dim=0
         )  # (B, C, H, W)
-
+        if self.do_resizing:
+            chips = self._resize_image(chips, chips.shape[2])
         order = self.metadata[self.modality]["band_order"]
         wavelengths = []
         for band in self.metadata[self.modality]["band_order"]:

rslearn/models/concatenate_features.py

@@ -0,0 +1,93 @@
+"""Concatenate feature map with features from input data."""
+
+from typing import Any
+
+import torch
+
+
+class ConcatenateFeatures(torch.nn.Module):
+    """Concatenate feature map with additional raw data inputs."""
+
+    def __init__(
+        self,
+        key: str,
+        in_channels: int | None = None,
+        conv_channels: int = 64,
+        out_channels: int | None = None,
+        num_conv_layers: int = 1,
+        kernel_size: int = 3,
+        final_relu: bool = False,
+    ):
+        """Create a new ConcatenateFeatures.
+
+        Args:
+            key: the key of the input_dict to concatenate.
+            in_channels: number of input channels of the additional features.
+            conv_channels: number of channels of the convolutional layers.
+            out_channels: number of output channels of the additional features.
+            num_conv_layers: number of convolutional layers to apply to the additional features.
+            kernel_size: kernel size of the convolutional layers.
+            final_relu: whether to apply a ReLU activation to the final output, default False.
+        """
+        super().__init__()
+        self.key = key
+
+        if num_conv_layers > 0:
+            if in_channels is None or out_channels is None:
+                raise ValueError(
+                    "in_channels and out_channels must be specified if num_conv_layers > 0"
+                )
+
+        conv_layers = []
+        for i in range(num_conv_layers):
+            conv_in = in_channels if i == 0 else conv_channels
+            conv_out = out_channels if i == num_conv_layers - 1 else conv_channels
+            conv_layers.append(
+                torch.nn.Conv2d(
+                    in_channels=conv_in,
+                    out_channels=conv_out,
+                    kernel_size=kernel_size,
+                    padding="same",
+                )
+            )
+            if i < num_conv_layers - 1 or final_relu:
+                conv_layers.append(torch.nn.ReLU(inplace=True))
+
+        self.conv_layers = torch.nn.Sequential(*conv_layers)
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> list[torch.Tensor]:
+        """Concatenate the feature map with the raw data inputs.
+
+        Args:
+            features: list of feature maps at different resolutions.
+            inputs: original inputs.
+
+        Returns:
+            concatenated feature maps.
+        """
+        if not features:
+            raise ValueError("Expected at least one feature map, got none.")
+
+        add_data = torch.stack([input_data[self.key] for input_data in inputs], dim=0)
+        add_features = self.conv_layers(add_data)
+
+        new_features: list[torch.Tensor] = []
+        for feature_map in features:
+            # Shape of feature map: BCHW
+            feat_h, feat_w = feature_map.shape[2], feature_map.shape[3]
+
+            resized_add_features = add_features
+            # Resize additional features to match each feature map size if needed
+            if add_features.shape[2] != feat_h or add_features.shape[3] != feat_w:
+                resized_add_features = torch.nn.functional.interpolate(
+                    add_features,
+                    size=(feat_h, feat_w),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+
+            new_features.append(torch.cat([feature_map, resized_add_features], dim=1))
+
+        return new_features

rslearn/models/croma.py

@@ -7,6 +7,7 @@ from enum import Enum
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 from upath import UPath
 
@@ -99,6 +100,7 @@ class Croma(torch.nn.Module):
         modality: CromaModality,
         pretrained_path: str | None = None,
         image_resolution: int = DEFAULT_IMAGE_RESOLUTION,
+        do_resizing: bool = False,
     ) -> None:
         """Instantiate a new Croma instance.
 
@@ -107,12 +109,21 @@ class Croma(torch.nn.Module):
             modality: the modalities to configure the model to accept.
             pretrained_path: the local path to the pretrained weights. Otherwise it is
                 downloaded and cached in temp directory.
-            image_resolution: the width and height of the input images.
+            image_resolution: the width and height of the input images passed to the model. if do_resizing is True, the image will be resized to this resolution.
+            do_resizing: Whether to resize the image to the input resolution.
         """
         super().__init__()
         self.size = size
         self.modality = modality
-        self.image_resolution = image_resolution
+        self.do_resizing = do_resizing
+        if not do_resizing:
+            self.image_resolution = image_resolution
+        else:
+            # With single pixel input, we always resample to the patch size.
+            if image_resolution == 1:
+                self.image_resolution = PATCH_SIZE
+            else:
+                self.image_resolution = DEFAULT_IMAGE_RESOLUTION
 
         # Cache the CROMA weights to a deterministic path in temporary directory if the
         # path is not provided by the user.
@@ -137,7 +148,16 @@ class Croma(torch.nn.Module):
             pretrained_path=pretrained_path,
             size=size.value,
             modality=modality.value,
-            image_resolution=image_resolution,
+            image_resolution=self.image_resolution,
+        )
+
+    def _resize_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Resize the image to the input resolution."""
+        return F.interpolate(
+            image,
+            size=(self.image_resolution, self.image_resolution),
+            mode="bilinear",
+            align_corners=False,
         )
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
@@ -151,8 +171,11 @@ class Croma(torch.nn.Module):
         sentinel2: torch.Tensor | None = None
         if self.modality in [CromaModality.BOTH, CromaModality.SENTINEL1]:
             sentinel1 = torch.stack([inp["sentinel1"] for inp in inputs], dim=0)
+            sentinel1 = self._resize_image(sentinel1) if self.do_resizing else sentinel1
         if self.modality in [CromaModality.BOTH, CromaModality.SENTINEL2]:
             sentinel2 = torch.stack([inp["sentinel2"] for inp in inputs], dim=0)
+            sentinel2 = self._resize_image(sentinel2) if self.do_resizing else sentinel2
+
         outputs = self.model(
             SAR_images=sentinel1,
             optical_images=sentinel2,