Merge pull request #11 from RadarML/dev/misc

Miscellaneous improvements & initial public release

Author: thetianshuhuang

README.md

@@ -1,5 +1,13 @@
 # NRDK: Neural Radar Development Kit
 
+[![GitHub](https://img.shields.io/github/v/release/RadarML/nrdk)](https://github.com/RadarML/nrdk)
+![MIT License](https://img.shields.io/badge/license-MIT-green)
+![Supports Python 3.12+](https://img.shields.io/badge/python-3.12%20%7C%203.13-blue)
+![Typed](https://img.shields.io/badge/types-typed-limegreen)
+[![bear-ified](https://raw.githubusercontent.com/beartype/beartype-assets/main/badge/bear-ified.svg)](https://beartype.readthedocs.io)
+[![CI](https://github.com/RadarML/nrdk/actions/workflows/ci.yml/badge.svg)](https://github.com/RadarML/nrdk/actions/workflows/ci.yml)
+![GitHub issues](https://img.shields.io/github/issues/RadarML/nrdk)
+
 > [!IMPORTANT]
 > See our [documentation site](https://radarml.github.io/nrdk) for more details!
 

docs/grt/index.md

@@ -30,7 +30,7 @@ example-project/
 
         Don't forget to change the `name`, `authors`, and `description`!
 
-2. Set up the `nrdk` dependency.
+2. Set up the `nrdk` dependency (`nrdk[roverd] >= 0.1.5`).
 
     !!! warning "Required Extras"
 
@@ -71,3 +71,9 @@ The GRT template includes reference training scripts which can be used for high
     ```python title="grt/train_minimal.py"
     --8<-- "grt/train_minimal.py"
     ```
+
+## Evaluation Script
+
+::: evaluate.evaluate
+    options:
+        show_root_heading: false

docs/index.md

@@ -1,5 +1,13 @@
 # NRDK: Neural Radar Development Kit
 
+[![GitHub](https://img.shields.io/github/v/release/RadarML/nrdk)](https://github.com/RadarML/nrdk)
+![MIT License](https://img.shields.io/badge/license-MIT-green)
+![Supports Python 3.12+](https://img.shields.io/badge/python-3.12%20%7C%203.13-blue)
+![Typed](https://img.shields.io/badge/types-typed-limegreen)
+[![bear-ified](https://raw.githubusercontent.com/beartype/beartype-assets/main/badge/bear-ified.svg)](https://beartype.readthedocs.io)
+[![CI](https://github.com/RadarML/nrdk/actions/workflows/ci.yml/badge.svg)](https://github.com/RadarML/nrdk/actions/workflows/ci.yml)
+![GitHub issues](https://img.shields.io/github/issues/RadarML/nrdk)
+
 The **Neural Radar Development Kit** (NRDK) is an open-source and MIT-licensed Python library and framework for developing, training, and evaluating machine learning models on radar spectrum and multimodal sensor data.
 
 Built around typed, high modular interfaces, the NRDK is designed to reduce the barrier of entry to learning on spectrum via out-of-the-box reference implementations for [red-rover](https://radarml.github.io/red-rover/) data and the [I/Q-1M Dataset](https://radarml.github.io/red-rover/iq1m/), while also providing an easy path towards customization and extensions for other radar and data collection systems.

grt/evaluate.py

@@ -10,6 +10,8 @@
 import numpy as np
 import torch
 import tyro
+import wadler_lindig as wl
+from abstract_dataloader import spec
 from omegaconf import DictConfig
 from roverd.channels.utils import Prefetch
 from roverd.sensors import DynamicSensor
@@ -18,6 +20,20 @@
 from nrdk.framework import Result
 
 
+class _DatasetMeta(spec.Dataset):
+    def __init__(
+        self, dataset: spec.Dataset[dict[str, Any]], meta: Any
+    ) -> None:
+        self.dataset = dataset
+        self.meta = meta
+
+    def __getitem__(self, index: int | np.integer) -> dict[str, Any]:
+        return {"meta": self.meta, **self.dataset[index]}
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+
 def _get_dataloaders(
     cfg: DictConfig, data_root: str, transforms: Any,
     traces: list[str] | None = None, filter: str | None = None,
@@ -36,19 +52,28 @@ def _get_dataloaders(
                 os.path.relpath(t, cfg["meta"]["dataset"])
                 for t in hydra.utils.instantiate(
                     cfg["datamodule"]["traces"]["test"])]
+
+        _unfiltered = traces
         if filter is not None:
             traces = [t for t in traces if re.match(filter, t)]
+        if len(traces) == 0:
+            raise ValueError(
+                f"No traces match the filter {filter}:\n"
+                f"{wl.pprint(_unfiltered)}")
 
         def construct(t: str) -> torch.utils.data.DataLoader:
-            dataset = dataset_constructor(paths=[t])
+            dataset = _DatasetMeta(
+                dataset_constructor(paths=[t]),
+                meta={"train": False, "split": "test"})
+
             return datamodule.dataloader(dataset, mode="test")
 
         return {
             t: partial(construct, os.path.join(data_root, t)) for t in traces}
 
 
 def evaluate(
-    path: str, /, sample: int | None = None,
+    path: str, /, output: str | None = None, sample: int | None = None,
     traces: list[str] | None = None, filter: str | None = None,
     data_root: str | None = None,
     device: str = "cuda:0",
@@ -80,6 +105,7 @@ def evaluate(
 
     Args:
         path: path to results directory.
+        output: if specified, write results to this directory instead.
         sample: number of samples to evaluate.
         traces: explicit list of traces to evaluate.
         filter: evaluate all traces matching this regex.
@@ -90,11 +116,16 @@ def evaluate(
         workers: number of workers for data loading.
         prefetch: number of batches to prefetch per worker.
     """
+    torch.set_float32_matmul_precision('high')
+
     result = Result(path)
     cfg = result.config()
     if sample is not None:
         cfg["datamodule"]["subsample"]["test"] = sample
 
+    if output is None:
+        output = os.path.join(path, "eval")
+
     if data_root is None:
         data_root = cfg["meta"]["dataset"]
         if data_root is None:
@@ -107,6 +138,7 @@ def evaluate(
     cfg["datamodule"]["batch_size"] = batch
     cfg["datamodule"]["num_workers"] = workers
     cfg["datamodule"]["prefetch_factor"] = prefetch
+    cfg["lightningmodule"]["compile"] = False
 
     transforms = hydra.utils.instantiate(cfg["transforms"])
     lightningmodule = hydra.utils.instantiate(
@@ -120,7 +152,7 @@ def evaluate(
     def collect_metadata(y_true):
         return {
             f"meta/{k}/ts": getattr(v, "timestamps")
-            for k, v in y_true.items()
+            for k, v in y_true.items() if hasattr(v, "timestamps")
         }
 
     for trace, dl_constructor in dataloaders.items():
@@ -131,8 +163,7 @@ def collect_metadata(y_true):
             total=len(dataloader), desc=trace)
 
         output_container = DynamicSensor(
-            os.path.join(result.path, "eval", trace),
-            create=True, exist_ok=True)
+            os.path.join(output, trace), create=True, exist_ok=True)
         metrics = []
         outputs = {}
         for batch_metrics, vis in eval_stream:
@@ -160,7 +191,7 @@ def collect_metadata(y_true):
             k: np.concatenate([m[k] for m in metrics], axis=0)
             for k in metrics[0]}
         np.savez_compressed(
-            os.path.join(result.path, "eval", trace, "metrics.npz"),
+            os.path.join(output, trace, "metrics.npz"),
             **metrics, allow_pickle=False)
 
         output_container.create("ts", meta={

grt/pyproject.toml

@@ -8,7 +8,7 @@ description = "GRT/NRDK Reference Implementation"
 requires-python = ">=3.12"
 
 dependencies = [
-    "nrdk[roverd] >= 0.1.1",
+    "nrdk[roverd] >= 0.1.5",
     "pyyaml >= 5.0.0",
     "hydra-core >= 1.3.0",
     "lightning >= 2.5.5",

grt/train.py

@@ -1,53 +1,88 @@
-"""GRT reference implementation training script."""
+"""GRT Reference implementation training script."""
 
 import logging
 import os
+from collections.abc import Mapping
 from time import perf_counter
 from typing import Any
 
 import hydra
 import torch
 import yaml
 from lightning.pytorch import callbacks
+from omegaconf import DictConfig
+from rich.logging import RichHandler
+
+from nrdk.framework import Result
+
+logger = logging.getLogger("train")
+
+def _configure_logging(cfg: DictConfig) -> None:
+    log_level = cfg.meta.get("verbose", logging.INFO)
+    root = logging.getLogger()
+    root.setLevel(log_level)
+    root.handlers.clear()
+
+    rich_handler = RichHandler(markup=True)
+    rich_handler.setFormatter(logging.Formatter(
+        "[orange1]%(name)s:[/orange1] %(message)s"))
+    root.addHandler(rich_handler)
+
+    logger.debug(f"Configured with log level: {log_level}")
+
+
+def _load_weights(lightningmodule, path: str, rename: Mapping = {}) -> None:
+    weights = Result(path).best if os.path.isdir(path) else path
+    lightningmodule.load_weights(weights, rename=rename)
+
+
+def _get_best(trainer) -> dict[str, Any]:
+    for callback in trainer.callbacks:
+        if isinstance(callback, callbacks.ModelCheckpoint):
+            return {
+                "best_k": {
+                    os.path.basename(k): v.item()
+                    for k, v in callback.best_k_models.items()},
+                "best": os.path.basename(callback.best_model_path)
+            }
+    return {}
 
 
 @hydra.main(version_base=None, config_path="./config", config_name="default")
-def train(cfg):
+def train(cfg: DictConfig) -> None:
     """Train a model using the GRT reference implementation."""
     torch.set_float32_matmul_precision('high')
+    _configure_logging(cfg)
+
+    if cfg["meta"]["name"] is None or cfg["meta"]["version"] is None:
+        logger.error("Must set `meta.name` and `meta.version` in the config.")
+        return
 
     def _inst(path, *args, **kwargs):
         return hydra.utils.instantiate(
             cfg[path], _convert_="all", *args, **kwargs)
 
-    if cfg["meta"]["name"] is None or cfg["meta"]["version"] is None:
-        logging.error("Must set `meta.name` and `meta.version` in the config.")
-        return
-
     transforms = _inst("transforms")
     datamodule = _inst("datamodule", transforms=transforms)
     lightningmodule = _inst("lightningmodule", transforms=transforms)
     trainer = _inst("trainer")
-
     if "base" in cfg:
-        lightningmodule.load_weights(
-            cfg['base']['path'], rename=cfg['base'].get('rename', {}))
+        _load_weights(lightningmodule, **cfg['base'])
 
     start = perf_counter()
+    logger.info(
+        f"Start training @ {cfg["meta"]["results"]}/{cfg["meta"]["name"]}/"
+        f"{cfg["meta"]["version"]} [t={start:.3f}]")
     trainer.fit(
         model=lightningmodule, datamodule=datamodule,
         ckpt_path=cfg['meta']['resume'])
     duration = perf_counter() - start
+    logger.info(
+        f"Training completed in {duration / 60 / 60:.2f}h (={duration:.3f}s).")
 
     meta: dict[str, Any] = {"duration": duration}
-    for callback in trainer.callbacks:
-        if isinstance(callback, callbacks.ModelCheckpoint):
-            meta["best_k"] = {
-                os.path.basename(k): v.item()
-                for k, v in callback.best_k_models.items()}
-            meta["best"] = os.path.basename(callback.best_model_path)
-            break
-
+    meta.update(_get_best(trainer))
+    logger.info(f"Best checkpoint: {meta.get('best')}")
     meta_path = os.path.join(trainer.logger.log_dir, "checkpoints.yaml")
     with open(meta_path, 'w') as f:
         yaml.dump(meta, f, sort_keys=False)

grt/uv.lock

@@ -4,7 +4,7 @@ repo_name: RadarML/nrdk
 nav:
   - nrdk:
     - Overview: index.md
-    - Architecture: design.md
+    - Software Architecture: design.md
     - Core Modules:
       - nrdk.config: nrdk/config.md
       - nrdk.framework: nrdk/framework.md
@@ -14,7 +14,7 @@ nav:
       - nrdk.objectives: nrdk/objectives.md
       - nrdk.roverd: nrdk/roverd.md
       - nrdk.vis: nrdk/vis.md
-    - CLI Tools: cli.md
+  - cli: cli.md
   - nrdk.tss:
     - Overview: tss/index.md
     - High Level API: tss/api.md

mkdocs.yml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "nrdk"
-version = "0.1.4"
+version = "0.1.5"
 authors = [
   { name="Tianshu Huang", email="tianshu2@andrew.cmu.edu" },
 ]

pyproject.toml

@@ -4,7 +4,6 @@
 import os
 import re
 import threading
-import warnings
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from functools import cache
 from typing import Any, Callable, Generic, TypeVar, cast
@@ -91,10 +90,12 @@ def __init__(
     ) -> None:
         super().__init__()
 
+        self._log = logging.getLogger("NRDKLightningModule")
+
         if compile:
             jt_disable = os.environ.get("JAXTYPING_DISABLE", "0").lower()
             if jt_disable not in ("1", "true"):
-                warnings.warn(
+                self._log.warning(
                     "torch.compile is currently incompatible with jaxtyping; "
                     "if you see type errors, set the environment variable "
                     "`JAXTYPING_DISABLE=1` to disable jaxtyping checks.")
@@ -108,8 +109,6 @@ def __init__(
         self.vis_interval = vis_interval
         self.vis_samples = vis_samples
 
-        self._log = logging.getLogger(self.__class__.__name__)
-
     @torch.compiler.disable
     def load_weights(
         self, path: str, rename: Sequence[Mapping[str, str | None]] = []
@@ -153,9 +152,14 @@ def load_weights(
         if "model" in weights:
             weights = weights["model"]
 
-        weights = {
-            k[6:] if k.startswith("model.")
-            else k: v for k, v in weights.items()}
+        def _strip_prefix(k):
+            if k.startswith("model."):
+                k = k[6:]
+            if k.startswith("_orig_mod."):
+                k = k[10:]
+            return k
+
+        weights = {_strip_prefix(k): v for k, v in weights.items()}
         for pattern in rename:
             pat, sub = next(iter(pattern.items()))
             if sub is None:
@@ -191,7 +195,7 @@ def _make_log(
 
         if len(images) > 0:
             if not isinstance(self.logger, LoggerWithImages):
-                warnings.warn(
+                self._log.warning(
                     "Tried to log visualizations, but the logger does not "
                     "implement the `LoggerWithImages` interface.")
             else:
@@ -222,6 +226,7 @@ def log_visualizations(
             y_pred: model output values.
             split: train/val split to put in the output path.
         """
+        self._log.debug(f"Logging visualizations @ i={self.global_step}")
         y_true, y_pred = optree.tree_map(
             lambda x: x[:self.vis_samples].cpu().detach(),
             (y_true, y_pred))  # type: ignore