rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/train/lightning_module.py

@@ -1,5 +1,6 @@
 """Default LightningModule for rslearn."""
 
+import json
 import os
 from typing import Any
 
@@ -7,12 +8,17 @@ import lightning as L
 import torch
 from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
 from PIL import Image
-from torch.optim import AdamW
-from torch.optim.lr_scheduler import ReduceLROnPlateau
 from upath import UPath
 
+from rslearn.log_utils import get_logger
+
+from .model_context import ModelContext, ModelOutput
+from .optimizer import AdamW, OptimizerFactory
+from .scheduler import PlateauScheduler, SchedulerFactory
 from .tasks import Task
 
+logger = get_logger(__name__)
+
 
 class RestoreConfig:
     """Configuration for restoring model parameters.
@@ -36,7 +42,7 @@ class RestoreConfig:
             restore_path_options: additional options for the restore_path to pass to
                 fsspec.
             selector: path in the torch dict containing the model parameters.
-            ignore_prefixes: prefixes to restore.
+            ignore_prefixes: prefixes to ignore from the state dict.
             remap_prefixes: list of (old_prefix, new_prefix) to rename parameters
                 starting with old_prefix to start with new_prefix instead.
         """
@@ -47,9 +53,9 @@ class RestoreConfig:
 
     def get_state_dict(self) -> dict[str, Any]:
         """Returns the state dict configured in this RestoreConfig."""
-        print(f"loading state dict from {self.restore_path}")
+        logger.info(f"loading state dict from {self.restore_path}")
         with self.restore_path.open("rb") as f:
-            state_dict = torch.load(f)
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
         for k in self.selector:
             state_dict = state_dict[k]
 
@@ -82,48 +88,71 @@ class RslearnLightningModule(L.LightningModule):
         self,
         model: torch.nn.Module,
         task: Task,
+        optimizer: OptimizerFactory | None = None,
+        scheduler: SchedulerFactory | None = None,
+        visualize_dir: str | None = None,
+        metrics_file: str | None = None,
+        restore_config: RestoreConfig | None = None,
+        print_parameters: bool = False,
+        print_model: bool = False,
+        # Deprecated options.
         lr: float = 1e-3,
         plateau: bool = False,
         plateau_factor: float = 0.1,
         plateau_patience: int = 10,
         plateau_min_lr: float = 0,
         plateau_cooldown: int = 0,
-        visualize_dir: str | None = None,
-        restore_config: RestoreConfig | None = None,
-        print_parameters: bool = False,
-        print_model: bool = False,
     ):
         """Initialize a new RslearnLightningModule.
 
         Args:
             model: the model
             task: the task to train on
-            lr: the initial learning rate
-            plateau: whether to enable plateau scheduler (default false)
-            plateau_factor: on plateau, factor to multiply learning rate by
-            plateau_patience: number of iterations with no improvement in val loss
-                before reducing learning rate
-            plateau_min_lr: minimum learning rate to reduce to
-            plateau_cooldown: number of iterations after reducing learning rate before
-                resetting plateau scheduler
+            optimizer: the optimizer factory.
+            scheduler: the learning rate scheduler factory.
             visualize_dir: during validation or testing, output visualizations to this
                 directory
+            metrics_file: file to save metrics to
             restore_config: specification of configuration to restore parameters from
                 a non-Lightning checkpoint.
             print_parameters: whether to print the list of model parameters after model
                 initialization
             print_model: whether to print the model after model initialization
+            lr: deprecated.
+            plateau: deprecated.
+            plateau_factor: deprecated.
+            plateau_patience: deprecated.
+            plateau_min_lr: deprecated.
+            plateau_cooldown: deprecated.
         """
         super().__init__()
         self.model = model
         self.task = task
-        self.lr = lr
-        self.plateau = plateau
-        self.plateau_factor = plateau_factor
-        self.plateau_patience = plateau_patience
-        self.plateau_min_lr = plateau_min_lr
-        self.plateau_cooldown = plateau_cooldown
         self.visualize_dir = visualize_dir
+        self.metrics_file = metrics_file
+        self.restore_config = restore_config
+
+        self.scheduler_factory: SchedulerFactory | None = None
+        if scheduler:
+            self.scheduler_factory = scheduler
+        elif plateau:
+            logger.warning(
+                "The plateau argument to RslearnLightningModule is deprecated and will be removed in a future version"
+            )
+            self.scheduler_factory = PlateauScheduler(
+                factor=plateau_factor,
+                patience=plateau_patience,
+                min_lr=plateau_min_lr,
+                cooldown=plateau_cooldown,
+            )
+
+        if optimizer:
+            self.optimizer_factory = optimizer
+        else:
+            logger.warning(
+                "Defaulting the optimizer to AdamW since an OptimizerFactory was not provided. In a future version, the optimizer will be a required argument."
+            )
+            self.optimizer_factory = AdamW(lr=lr)
 
         if print_parameters:
             for name, param in self.named_parameters():
@@ -132,23 +161,26 @@ class RslearnLightningModule(L.LightningModule):
         if print_model:
             print(self.model)
 
-        if restore_config:
-            state_dict = restore_config.get_state_dict()
-            missing_keys, unexpected_keys = self.model.load_state_dict(
-                state_dict, strict=False
-            )
-            if missing_keys or unexpected_keys:
-                print(
-                    f"warning: restore yielded missing_keys={missing_keys} and unexpected_keys={unexpected_keys}"
-                )
-
         self.epochs = 0
 
         metrics = self.task.get_metrics()
         self.val_metrics = metrics.clone(prefix="val_")
         self.test_metrics = metrics.clone(prefix="test_")
 
-        self.schedulers = {}
+        self.schedulers: dict = {}
+
+    def on_fit_start(self) -> None:
+        """Called when the fit begins."""
+        # Only restore if doing a fresh fit.
+        if self.trainer.ckpt_path is None and self.restore_config:
+            state_dict = self.restore_config.get_state_dict()
+            missing_keys, unexpected_keys = self.model.load_state_dict(
+                state_dict, strict=False
+            )
+            if missing_keys or unexpected_keys:
+                logger.warning(
+                    f"restore yielded missing_keys={missing_keys} and unexpected_keys={unexpected_keys}"
+                )
 
     def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
         """Initialize the optimizer and learning rate scheduler.
@@ -156,27 +188,37 @@ class RslearnLightningModule(L.LightningModule):
         Returns:
             Optimizer and learning rate scheduler.
         """
-        params = [p for p in self.parameters() if p.requires_grad]
-        optimizer = AdamW(params, lr=self.lr)
+        optimizer = self.optimizer_factory.build(self)
         d = dict(
             optimizer=optimizer,
         )
-        if self.plateau:
-            scheduler = ReduceLROnPlateau(
-                optimizer,
-                factor=self.plateau_factor,
-                patience=self.plateau_patience,
-                min_lr=self.plateau_min_lr,
-                cooldown=self.plateau_cooldown,
-            )
+        if self.scheduler_factory is not None:
+            scheduler = self.scheduler_factory.build(optimizer)
             d["lr_scheduler"] = {
                 "scheduler": scheduler,
                 "monitor": "train_loss",
                 "interval": "epoch",
             }
-            self.schedulers["plateau"] = scheduler
+            self.schedulers["scheduler"] = scheduler
         return d
 
+    def on_train_epoch_start(self) -> None:
+        """If we are in a multi-dataset distributed strategy, set the epoch."""
+        try:
+            self.trainer.train_dataloader.batch_sampler.set_epoch(self.current_epoch)
+        except AttributeError:
+            # Fail silently for single-dataset case, which is okay
+            pass
+
+    def on_test_epoch_end(self) -> None:
+        """Optionally save the test metrics to a file."""
+        if self.metrics_file:
+            with open(self.metrics_file, "w") as f:
+                metrics = self.test_metrics.compute()
+                metrics_dict = {k: v.item() for k, v in metrics.items()}
+                json.dump(metrics_dict, f, indent=4)
+                logger.info(f"Saved metrics to {self.metrics_file}")
+
     def training_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
     ) -> torch.Tensor:
@@ -190,9 +232,16 @@ class RslearnLightningModule(L.LightningModule):
         Returns:
             The loss tensor.
         """
-        inputs, targets, _ = batch
+        inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        _, loss_dict = self(inputs, targets)
+        model_outputs = self(context, targets)
+        self.on_train_forward(context, targets, model_outputs)
+
+        loss_dict = model_outputs.loss_dict
         train_loss = sum(loss_dict.values())
         self.log_dict(
             {"train_" + k: v for k, v in loss_dict.items()},
@@ -200,6 +249,7 @@ class RslearnLightningModule(L.LightningModule):
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            sync_dist=True,
         )
         self.log(
             "train_loss",
@@ -207,6 +257,7 @@ class RslearnLightningModule(L.LightningModule):
             batch_size=batch_size,
             on_step=False,
             on_epoch=True,
+            sync_dist=True,
         )
         return train_loss
 
@@ -220,15 +271,24 @@ class RslearnLightningModule(L.LightningModule):
             batch_idx: Integer displaying index of this batch.
             dataloader_idx: Index of the current dataloader.
         """
-        inputs, targets, _ = batch
+        inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        outputs, loss_dict = self(inputs, targets)
+        model_outputs = self(context, targets)
+        self.on_val_forward(context, targets, model_outputs)
+
+        loss_dict = model_outputs.loss_dict
+        outputs = model_outputs.outputs
         val_loss = sum(loss_dict.values())
         self.log_dict(
             {"val_" + k: v for k, v in loss_dict.items()},
             batch_size=batch_size,
             on_step=False,
             on_epoch=True,
+            sync_dist=True,
         )
         self.log(
             "val_loss",
@@ -237,9 +297,12 @@ class RslearnLightningModule(L.LightningModule):
             prog_bar=True,
             on_step=False,
             on_epoch=True,
+            sync_dist=True,
         )
         self.val_metrics.update(outputs, targets)
-        self.log_dict(self.val_metrics, batch_size=batch_size, on_epoch=True)
+        self.log_dict(
+            self.val_metrics, batch_size=batch_size, on_epoch=True, sync_dist=True
+        )
 
     def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
         """Compute the test loss and additional metrics.
@@ -250,20 +313,36 @@ class RslearnLightningModule(L.LightningModule):
             dataloader_idx: Index of the current dataloader.
         """
         inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        outputs, loss_dict = self(inputs, targets)
+        model_outputs = self(context, targets)
+        self.on_test_forward(context, targets, model_outputs)
+
+        loss_dict = model_outputs.loss_dict
+        outputs = model_outputs.outputs
         test_loss = sum(loss_dict.values())
         self.log_dict(
             {"test_" + k: v for k, v in loss_dict.items()},
             batch_size=batch_size,
             on_step=False,
             on_epoch=True,
+            sync_dist=True,
         )
         self.log(
-            "test_loss", test_loss, batch_size=batch_size, on_step=False, on_epoch=True
+            "test_loss",
+            test_loss,
+            batch_size=batch_size,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
         )
         self.test_metrics.update(outputs, targets)
-        self.log_dict(self.test_metrics, batch_size=batch_size, on_epoch=True)
+        self.log_dict(
+            self.test_metrics, batch_size=batch_size, on_epoch=True, sync_dist=True
+        )
 
         if self.visualize_dir:
             for idx, (inp, target, output, metadata) in enumerate(
@@ -273,13 +352,13 @@ class RslearnLightningModule(L.LightningModule):
                 for image_suffix, image in images.items():
                     out_fname = os.path.join(
                         self.visualize_dir,
-                        f'{metadata["window_name"]}_{metadata["bounds"][0]}_{metadata["bounds"][1]}_{image_suffix}.png',
+                        f"{metadata['window_name']}_{metadata['bounds'][0]}_{metadata['bounds'][1]}_{image_suffix}.png",
                     )
                     Image.fromarray(image).save(out_fname)
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> torch.Tensor:
+    ) -> ModelOutput:
         """Compute the predicted class probabilities.
 
         Args:
@@ -290,18 +369,69 @@ class RslearnLightningModule(L.LightningModule):
         Returns:
             Output predicted probabilities.
         """
-        inputs, _, _ = batch
-        outputs, _ = self(inputs)
-        return outputs
+        inputs, _, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
+        model_outputs = self(context)
+        return model_outputs
 
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
+    def forward(
+        self, context: ModelContext, targets: list[dict[str, Any]] | None = None
+    ) -> ModelOutput:
         """Forward pass of the model.
 
         Args:
-            args: Arguments to pass to model.
-            kwargs: Keyword arguments to pass to model.
+            context: the model context.
+            targets: the target dicts.
 
         Returns:
             Output of the model.
         """
-        return self.model(*args, **kwargs)
+        return self.model(context, targets)
+
+    def on_train_forward(
+        self,
+        context: ModelContext,
+        targets: list[dict[str, Any]],
+        model_outputs: ModelOutput,
+    ) -> None:
+        """Hook to run after the forward pass of the model during training.
+
+        Args:
+            context: The model context.
+            targets: The target batch.
+            model_outputs: The output of the model.
+        """
+        pass
+
+    def on_val_forward(
+        self,
+        context: ModelContext,
+        targets: list[dict[str, Any]],
+        model_outputs: ModelOutput,
+    ) -> None:
+        """Hook to run after the forward pass of the model during validation.
+
+        Args:
+            context: The model context.
+            targets: The target batch.
+            model_outputs: The output of the model.
+        """
+        pass
+
+    def on_test_forward(
+        self,
+        context: ModelContext,
+        targets: list[dict[str, Any]],
+        model_outputs: ModelOutput,
+    ) -> None:
+        """Hook to run after the forward pass of the model during testing.
+
+        Args:
+            context: The model context.
+            targets: The target batch.
+            model_outputs: The output of the model.
+        """
+        pass

rslearn/train/model_context.py

@@ -0,0 +1,88 @@
+"""Data classes to provide various context to models."""
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+import torch
+
+from rslearn.utils.geometry import PixelBounds, Projection
+
+
+@dataclass
+class RasterImage:
+    """A raster image is a torch.tensor containing the images and their associated timestamps."""
+
+    # image is a 4D CTHW tensor
+    image: torch.Tensor
+    # if timestamps is not None, len(timestamps) must match the T dimension of the tensor
+    timestamps: list[tuple[datetime, datetime]] | None = None
+
+    @property
+    def shape(self) -> torch.Size:
+        """The shape of the image."""
+        return self.image.shape
+
+    def dim(self) -> int:
+        """The dim of the image."""
+        return self.image.dim()
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """The image dtype."""
+        return self.image.dtype
+
+    def single_ts_to_chw_tensor(self) -> torch.Tensor:
+        """Single timestep models expect single timestep inputs.
+
+        This function (1) checks this raster image only has 1 timestep and
+        (2) returns the tensor for that (single) timestep (going from CTHW to CHW).
+        """
+        if self.image.shape[1] != 1:
+            raise ValueError(f"Expected a single timestep, got {self.image.shape[1]}")
+        return self.image[:, 0]
+
+
+@dataclass
+class SampleMetadata:
+    """Metadata pertaining to an example."""
+
+    window_group: str
+    window_name: str
+    window_bounds: PixelBounds
+    patch_bounds: PixelBounds
+    patch_idx: int
+    num_patches_in_window: int
+    time_range: tuple[datetime, datetime] | None
+    projection: Projection
+
+    # Task name to differentiate different tasks.
+    dataset_source: str | None
+
+
+@dataclass
+class ModelContext:
+    """Context to pass to all model components."""
+
+    # One input dict per example in the batch.
+    inputs: list[dict[str, torch.Tensor | RasterImage]]
+    # One SampleMetadata per example in the batch.
+    metadatas: list[SampleMetadata]
+    # Arbitrary dict that components can add to.
+    context_dict: dict[str, Any] = field(default_factory=lambda: {})
+
+
+@dataclass
+class ModelOutput:
+    """The output from the Predictor.
+
+    Args:
+        outputs: output compatible with the configured Task.
+        loss_dict: map from loss names to scalar tensors.
+        metadata: arbitrary dict that can be used to store other outputs.
+    """
+
+    outputs: Iterable[Any]
+    loss_dict: dict[str, torch.Tensor]
+    metadata: dict[str, Any] = field(default_factory=lambda: {})

rslearn/train/optimizer.py

@@ -0,0 +1,31 @@
+"""Optimizers for rslearn."""
+
+from dataclasses import asdict, dataclass
+
+import lightning as L
+import torch.optim
+from torch.optim import Optimizer
+
+
+class OptimizerFactory:
+    """A factory class that initializes the optimizer given the LightningModule."""
+
+    def build(self, lm: L.LightningModule) -> Optimizer:
+        """Build the optimizer configured by this factory class."""
+        raise NotImplementedError
+
+
+@dataclass
+class AdamW(OptimizerFactory):
+    """Factory for AdamW optimzier."""
+
+    lr: float = 0.001
+    betas: tuple[float, float] = (0.9, 0.999)
+    eps: float | None = None
+    weight_decay: float | None = None
+
+    def build(self, lm: L.LightningModule) -> Optimizer:
+        """Build the AdamW optimizer."""
+        params = [p for p in lm.parameters() if p.requires_grad]
+        kwargs = {k: v for k, v in asdict(self).items() if v is not None}
+        return torch.optim.AdamW(params, **kwargs)