Merge pull request #16 from RadarML/dev/export

thetianshuhuang · web-flow · commit a6358f7e257d · 2025-12-02T10:08:50.000-05:00
Checkpoint Exporting &amp; Bugfixes
diff --git a/docs/grt/index.md b/docs/grt/index.md
@@ -10,6 +10,10 @@ The GRT reference implementation uses a [hydra](https://hydra.cc/docs/intro/) +
     2. Create a virtual environment in `nrdk/grt` with `uv sync`.
     3. Run with `uv run train.py`; see the hydra config files in `nrdk/grt/config/` for options.
 
+!!! info "Pre-Trained Checkpoints"
+
+    Pre-trained model checkpoints for the GRT reference implementation on the [I/Q-1M dataset](https://radarml.github.io/red-rover/iq1m/) can also be found [here](https://radarml.github.io/red-rover/iq1m/osn/#download-from-osn).
+
 ## Quick Start
 
 1. Create a new repository, and copy the contents of the `grt/` directory:
diff --git a/grt/config/model/decoder/lidar2d.yaml b/grt/config/model/decoder/lidar2d.yaml
@@ -14,7 +14,7 @@ occ2d:
   d_model: ${size.d_model}
   num_layers: ${size.dec_layers}
   shape: [1, 1, 1024, 256]
-  scale: [1.0, 1.0, 1.0, 1.0]
+  scale: [16.0, 16.0, 16.0, 16.0]
   w_min: 0.2
   patch: [1, 1, 16, 16]
   out_dim: 0
diff --git a/grt/config/model/decoder/lidar3d.yaml b/grt/config/model/decoder/lidar3d.yaml
@@ -14,7 +14,7 @@ occ3d:
   d_model: ${size.d_model}
   num_layers: 4
   shape: [1, 64, 128, 64]
-  scale: [1.0, 1.0, 1.0, 1.0]
+  scale: [16.0, 16.0, 16.0, 16.0]
   w_min: 0.2
   patch: [1, 8, 8, 8]
   out_dim: 0
diff --git a/grt/config/model/decoder/semseg.yaml b/grt/config/model/decoder/semseg.yaml
@@ -14,7 +14,7 @@ semseg:
   d_model: ${size.d_model}
   num_layers: 4
   shape: [1, 160, 160, 1]
-  scale: [1.0, 1.0, 0.666, 1.0]
+  scale: [16.0, 16.0, 10.666, 16.0]
   w_min: 0.2
   patch: [1, 5, 5, 1]
   out_dim: 8
diff --git a/grt/config/model/tokenizer/grt.yaml b/grt/config/model/tokenizer/grt.yaml
@@ -3,6 +3,6 @@ d_model: ${size.d_model}
 patch: [1, 2, 2, 8, 4]  # (time, doppler, elevation, azimuth, range)
 squeeze: [2, 3]  # elevation + azimuth
 n_channels: 2  # (amplitude, phase angle)
-scale: [1.0, 1.0, 1.0]
+scale: [16.0, 16.0, 16.0]
 w_min: 0.2
 positions: nd
diff --git a/grt/config/sensors/lidar.yaml b/grt/config/sensors/lidar.yaml
@@ -1,6 +1,6 @@
 lidar:
   _partial_: true
   _target_: roverd.sensors.OSLidarDepth
-  correction: auto
+  correction: null
   past: 0
   future: 0
diff --git a/grt/config/sensors/radar.yaml b/grt/config/sensors/radar.yaml
@@ -1,6 +1,6 @@
 radar:
   _partial_: true
   _target_: roverd.sensors.XWRRadar
-  correction: auto
+  correction: null
   past: 0
   future: 0
diff --git a/grt/config/sensors/semseg.yaml b/grt/config/sensors/semseg.yaml
@@ -1,6 +1,6 @@
 _camera:
   _partial_: true
   _target_: roverd.sensors.Semseg
-  correction: auto
+  correction: null
   past: 0
   future: 0
diff --git a/grt/config/size/small.yaml b/grt/config/size/small.yaml
@@ -1,5 +1,5 @@
 d_model: 512
-d_feedforward: 1024
+d_feedforward: 2048
 nhead: 8
 enc_layers: 4
 dec_layers: 4
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nrdk"
-version = "0.1.5"
+version = "0.1.6"
 authors = [
   { name="Tianshu Huang", email="tianshu2@andrew.cmu.edu" },
 ]
@@ -36,11 +36,12 @@ dependencies = [
     "scipy >= 1.10.0",
     "scipy-stubs",
     "pandas >= 1.0.0",
-    "pandas-stubs"
+    "pandas-stubs",
+    "hydra-core >= 1.3.0",
 ]
 
 [tool.uv.sources]
-roverd = { git = "ssh://git@github.com/radarml/red-rover.git", subdirectory = "format"}
+roverd = { git = "https://github.com/radarml/red-rover.git", subdirectory = "format"}
 
 [project.optional-dependencies]
 roverd = [
diff --git a/src/nrdk/_cli/export.py b/src/nrdk/_cli/export.py
@@ -5,17 +5,19 @@
 
 import torch
 import yaml
+from omegaconf import OmegaConf
 
 
 def cli_export(
-    path: str, /, output: str = "weights.pth"
+    path: str, /, output: str = "weights.pth", config: str = "model.yaml"
 ) -> None:
     """Export model weights from a full-service checkpoint.
 
     !!! info "Usage"
 
         Take the best checkpoint in `results/experiment/version`, and export
-        the model to `results/experiment/version/weights.pth`:
+        the model to `results/experiment/version/weights.pth` and config to
+        `results/experiment/version/model.yaml`:
         ```sh
         nrdk export results/experiment/version
         ```
@@ -28,17 +30,31 @@ def cli_export(
     - If the `path` points to a file, export that checkpoint.
     - If the `path` is a directory, the directory should have a
         `checkpoints.yaml` with a `best` key which specifies the best
-        checkpoint in a `checkpoints/` directory; the exported model is saved
-        relative to the `path`.
+        checkpoint in a `checkpoints/` directory; the exported model and config
+        are saved relative to the `path`.
+
+    !!! warning
+
+        If `path` is a file (i.e., a specific checkpoint), the model config
+        will not be exported (since its path not defined).
 
     Args:
         path: path to the checkpoint.
         output: path to save the exported weights.
+        config: path to save the exported model config.
     """
     if os.path.isdir(path):
-        output = os.path.join(path, output)
+        hydra_cfg = OmegaConf.load(os.path.join(path, ".hydra", "config.yaml"))
+        resolved = OmegaConf.to_container(hydra_cfg, resolve=True)
+        assert isinstance(resolved, dict)
+        # Save to config file
+        with open(os.path.join(path, config), 'w') as f:
+            yaml.dump({
+                "model": resolved["model"],
+                "transforms": resolved["transforms"],
+            }, f, default_flow_style=False, sort_keys=False)
 
-    if os.path.isdir(path):
+        output = os.path.join(path, output)
         try:
             with open(os.path.join(path, "checkpoints.yaml"), 'r') as f:
                 contents = yaml.safe_load(f)
@@ -52,7 +68,7 @@ def cli_export(
 
     print(f"Exporting: {path}")
     print(f"--> {output}")
-    pattern = re.compile(r"^model\.")
+    pattern = re.compile(r"^(model\.)(_orig_mod\.)?")
     state_dict = {
         pattern.sub("", k): v for k, v in contents['state_dict'].items()}
     torch.save(state_dict, output)
diff --git a/src/nrdk/config/__init__.py b/src/nrdk/config/__init__.py
@@ -3,6 +3,6 @@
 from jaxtyping import install_import_hook
 
 with install_import_hook("nrdk.config", "beartype.beartype"):
-    from .config import Singleton, expand
+    from .config import Singleton, expand, inst_from
 
-__all__ = ["expand", "Singleton"]
+__all__ = ["expand", "inst_from", "Singleton"]
diff --git a/src/nrdk/config/config.py b/src/nrdk/config/config.py
@@ -4,6 +4,9 @@
 from collections.abc import Mapping, Sequence
 from typing import Any
 
+import hydra
+import yaml
+
 Nested = Sequence[str] | Mapping[str, "Nested"]
 
 
@@ -88,3 +91,29 @@ def __call__(self, name: str) -> Any:
         name: example
     ```
 """
+
+
+def inst_from(
+    path: str, key: str | Sequence[str] | None
+) -> Any:
+    """Load and instantiate an object from an arbitrary configuration file.
+
+    Args:
+        path: path to the configuration yaml file.
+        key: key or sequence of keys to locate the object specification in the
+            configuration file. If `None`, the entire configuration is used.
+
+    Returns:
+        Instantiated object (via `hydra.utils.instantiate`).
+    """
+    with open(path) as f:
+        spec = yaml.safe_load(f)
+
+    if key is not None:
+        if isinstance(key, str):
+            spec = spec[key]
+        else:
+            for k in key:
+                spec = spec[k]
+    return hydra.utils.instantiate(spec)
+
diff --git a/src/nrdk/framework/__init__.py b/src/nrdk/framework/__init__.py
@@ -12,14 +12,16 @@
 from jaxtyping import install_import_hook
 
 with install_import_hook("nrdk.framework", "beartype.beartype"):
-    from .architectures import TokenizerEncoderDecoder
+    from .architectures import Output, TokenizerEncoderDecoder
     from .callbacks import PerformanceMonitor
     from .lightningmodule import NRDKLightningModule
+    from .load import load_model
     from .logging import LoggerWithImages, MLFlowLogger, TensorBoardLogger
     from .result import Result
 
 __all__ = [
-    "TokenizerEncoderDecoder",
+    "load_model",
+    "Output", "TokenizerEncoderDecoder",
     "PerformanceMonitor",
     "LoggerWithImages", "MLFlowLogger", "TensorBoardLogger",
     "NRDKLightningModule", "Result"]
diff --git a/src/nrdk/framework/architectures.py b/src/nrdk/framework/architectures.py
@@ -1,14 +1,39 @@
 """Generic model architectures."""
 
 from collections.abc import Mapping
-from typing import Any
+from dataclasses import dataclass
+from typing import Any, TypeVar
 
 from torch import nn
 
+TReg = TypeVar("TReg", bound=dict[str, Any])
+
+@dataclass
+class Output[Tstage, Toutput]:
+    """Module with "shortcut" outputs which are directly passed to the output.
+
+    Type Parameters:
+        - Tstage: type of the primary module output.
+        - Toutput: type of the regularization side channel; must be a dictionary
+            with string keys.
+
+    Attributes:
+        stage: output for this stage to pass to the next stage.
+        output: outputs to pass directly to the model output.
+    """
+
+    stage: Tstage
+    output: Toutput
+
 
 class TokenizerEncoderDecoder(nn.Module):
     """Generic architecture with a tokenizer, encoder, and decoder(s).
 
+    The `tokenizer` and `encoder` can optionally return a [`Output`][^.]
+    output, in which case the regularization outputs (`reg`) are returned
+    directly in the model output, with only the primary output (`out`) passed
+    to the next stage in the model.
+
     !!! info
 
         The user/caller is responsible for ensuring that the dataloader output
@@ -50,9 +75,19 @@ def forward(self, data: dict[str, Any]) -> dict[str, Any]:
                 f"input key `{self.key}`, but only the following are "
                 f"available\n:{list(data.keys())}")
 
+        outputs = {}
+
         x = data[self.key]
         tokens = self.tokenizer(x)
+        if isinstance(tokens, Output):
+            outputs.update(tokens.output)
+            tokens = tokens.stage
+
         encoded = self.encoder(tokens)
+        if isinstance(encoded, Output):
+            outputs.update(encoded.output)
+            encoded = encoded.stage
+
         decoded = {k: v(encoded) for k, v in self.decoder.items()}
 
         # 0 - batch; 1 - temporal
@@ -61,4 +96,4 @@ def forward(self, data: dict[str, Any]) -> dict[str, Any]:
                 k: v.squeeze(dim=tuple(range(2, v.ndim)))
                 for k, v in decoded.items()}
 
-        return decoded
+        return {**outputs, **decoded}
diff --git a/src/nrdk/framework/load.py b/src/nrdk/framework/load.py
@@ -0,0 +1,49 @@
+"""Load pre-trained model for inference."""
+
+import os
+
+import hydra
+import torch
+import yaml
+
+
+def load_model(
+    path: str, config: str | None = None, freeze: bool = True
+) -> torch.nn.Module:
+    """Load a pre-trained model from a given path.
+
+    If `path` is a directory, it should contain a `weights.pth` file and a
+    `config.yaml` file, i.e., as saved by
+    [`nrdk export`][nrdk._cli.cli_export].
+
+    Args:
+        path: path to the model checkpoint or directory.
+        config: if `path` is a file, provide the model configuration separately
+            here.
+        freeze: whether to freeze the model parameters.
+
+    Returns:
+        Loaded model.
+    """
+    if os.path.isdir(path):
+        if config is None:
+            config = os.path.join(path, "model.yaml")
+        path = os.path.join(path, "weights.pth")
+
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"Model weights file not found: {path}")
+    if config is None or not os.path.isfile(config):
+        raise FileNotFoundError(f"Model config file not found: {config}")
+
+    with open(config) as f:
+        model_spec = yaml.safe_load(f)
+
+    model = hydra.utils.instantiate(model_spec["model"])
+    weights = torch.load(path)
+    model.load_state_dict(weights, strict=True)
+
+    if freeze:
+        for param in model.parameters():
+            param.requires_grad = False
+
+    return model
diff --git a/src/nrdk/objectives/occupancy.py b/src/nrdk/objectives/occupancy.py
@@ -192,7 +192,7 @@ def render(
         if render_gt:
             rendered["occ3d_gt"] = y_true.occupancy.to(
                 torch.uint8).cpu().numpy()
-        return  rendered
+        return rendered
 
 
 class Occupancy2D(Objective[
diff --git a/src/nrdk/vis/utils.py b/src/nrdk/vis/utils.py
@@ -48,7 +48,8 @@ def tile_images(
         y_hat: optional secondary images to tile; if provided, each `y_hat`
             is placed under its corresponding `y_true`.
         cols: number of columns. If the number of images does not divide
-            `batch`, they are zero padded in the resulting image.
+            `batch`, they are zero padded in the resulting image unless there
+            are fewer images than `batch`.
         cmap: colormap to apply, either as the name of a matplotlib colormap
             or a list of colors.
         normalize: whether to normalize the input data. All data is normalized
@@ -75,7 +76,9 @@ def tile_images(
     else:
         images = torch.concatenate([y_true] + _y_hat, dim=1)
 
-    if images.shape[0] % cols != 0:
+    if images.shape[0] < cols:
+        cols = images.shape[0]
+    elif images.shape[0] % cols != 0:
         padding = torch.zeros(
             [cols - (images.shape[0] % cols), *images.shape[1:]],
             dtype=images.dtype, device=images.device)
diff --git a/update_gh_pages.sh b/update_gh_pages.sh
diff --git a/uv.lock b/uv.lock
