rslearn 0.0.15__py3-none-any.whl → 0.0.17__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,273 @@ 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
+
+            # Avoid error in case project_dir doesn't exist.
+            fnames = project_dir.iterdir() if project_dir.exists() else []
+
+            for option in fnames:
+                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)
+        elif c.trainer.logger:
+            logger.warning(
+                "Model management is enabled and logging should be off, but the model config specifies a logger. "
+                + "The logger should be removed from the model config, since it will not be automatically disabled."
+            )
+
+        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 +395,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 +415,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/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/olmoearth_pretrain/model.py

@@ -153,11 +153,8 @@ class OlmoEarth(torch.nn.Module):
         # Load the checkpoint.
         if not random_initialization:
             train_module_dir = checkpoint_upath / "model_and_optim"
-            if train_module_dir.exists():
-                load_model_and_optim_state(str(train_module_dir), model)
-                logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
-            else:
-                logger.info(f"could not find OlmoEarth encoder at {train_module_dir}")
+            load_model_and_optim_state(str(train_module_dir), model)
+            logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
 
         return model
 

rslearn/tile_stores/__init__.py

@@ -22,17 +22,6 @@ def load_tile_store(config: dict[str, Any], ds_path: UPath) -> TileStore:
     Returns:
         the TileStore
     """
-    if config is None:
-        tile_store = DefaultTileStore()
-        tile_store.set_dataset_path(ds_path)
-        return tile_store
-
-    # Backwards compatability.
-    if "name" in config and "root_dir" in config and config["name"] == "file":
-        tile_store = DefaultTileStore(config["root_dir"])
-        tile_store.set_dataset_path(ds_path)
-        return tile_store
-
     init_jsonargparse()
     parser = jsonargparse.ArgumentParser()
     parser.add_argument("--tile_store", type=TileStore)

rslearn/train/dataset.py

@@ -17,9 +17,7 @@ from rasterio.warp import Resampling
 import rslearn.train.transforms.transform
 from rslearn.config import (
     DType,
-    RasterFormatConfig,
-    RasterLayerConfig,
-    VectorLayerConfig,
+    LayerConfig,
 )
 from rslearn.dataset.dataset import Dataset
 from rslearn.dataset.window import Window, get_layer_and_group_from_dir_name
@@ -28,8 +26,6 @@ from rslearn.train.tasks import Task
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.mp import star_imap_unordered
-from rslearn.utils.raster_format import load_raster_format
-from rslearn.utils.vector_format import load_vector_format
 
 from .transforms import Sequential
 
@@ -185,7 +181,7 @@ def read_raster_layer_for_data_input(
     bounds: PixelBounds,
     layer_name: str,
     group_idx: int,
-    layer_config: RasterLayerConfig,
+    layer_config: LayerConfig,
     data_input: DataInput,
 ) -> torch.Tensor:
     """Read a raster layer for a DataInput.
@@ -246,9 +242,7 @@ def read_raster_layer_for_data_input(
         )
         if band_set.format is None:
             raise ValueError(f"No format specified for {layer_name}")
-        raster_format = load_raster_format(
-            RasterFormatConfig(band_set.format["name"], band_set.format)
-        )
+        raster_format = band_set.instantiate_raster_format()
         raster_dir = window.get_raster_dir(
             layer_name, band_set.bands, group_idx=group_idx
         )
@@ -349,7 +343,6 @@ def read_data_input(
         images: list[torch.Tensor] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            assert isinstance(layer_config, RasterLayerConfig)
             images.append(
                 read_raster_layer_for_data_input(
                     window, bounds, layer_name, group_idx, layer_config, data_input
@@ -363,8 +356,7 @@ def read_data_input(
         features: list[Feature] = []
         for layer_name, group_idx in layers_to_read:
             layer_config = dataset.layers[layer_name]
-            assert isinstance(layer_config, VectorLayerConfig)
-            vector_format = load_vector_format(layer_config.format)
+            vector_format = layer_config.instantiate_vector_format()
             layer_dir = window.get_layer_dir(layer_name, group_idx=group_idx)
             cur_features = vector_format.decode_vector(
                 layer_dir, window.projection, window.bounds