Merge pull request #20 from RadarML/dev/compile

Compile Overhaul

Author: thetianshuhuang

docs/grt/index.md

@@ -10,47 +10,47 @@ 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 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).
+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).
 
-    With a single GPU, these checkpoints can be reproduced with the following:
+With a single GPU, these checkpoints can be reproduced with the following:
 
-    === "Base Model"
-        The 3D polar occupancy base model is provided as the default configuration (i.e., `sensors=[radar,lidar]`, `transforms@transforms.sample=[radar,lidar3d]`, `objective=lidar3d`, `model/decoder=lidar3d`).
-        ```sh
-        uv run train.py meta.name=base meta.version=small size=small
-        ```
+=== "Base Model"
+    The 3D polar occupancy base model is provided as the default configuration (i.e., `sensors=[radar,lidar]`, `transforms@transforms.sample=[radar,lidar3d]`, `objective=lidar3d`, `model/decoder=lidar3d`).
+    ```sh
+    uv run train.py meta.name=base meta.version=small size=small
+    ```
 
-    === "2D Occupancy"
-        ```sh
-        uv run train.py meta.name=occ2d meta.version=small size=small \
-            +base=occ3d_to_occ2d \
-            sensors=[radar,lidar] \
-            transforms@transforms.sample=[radar,lidar2d] \
-            objective=lidar2d \
-            model/decoder=lidar2d
-        ```
+=== "2D Occupancy"
+    ```sh
+    uv run train.py meta.name=occ2d meta.version=small size=small \
+        +base=occ3d_to_occ2d \
+        sensors=[radar,lidar] \
+        transforms@transforms.sample=[radar,lidar2d] \
+        objective=lidar2d \
+        model/decoder=lidar2d
+    ```
 
-    === "Semantic Segmentation"
-        ```sh
-        uv run train.py meta.name=semseg meta.version=small size=small \
-            +base=occ3d_to_semseg \
-            sensors=[radar,semseg] \
-            transforms@transforms.sample=[radar,semseg] \
-            objective=semseg \
-            model/decoder=semseg
-        ```
+=== "Semantic Segmentation"
+    ```sh
+    uv run train.py meta.name=semseg meta.version=small size=small \
+        +base=occ3d_to_semseg \
+        sensors=[radar,semseg] \
+        transforms@transforms.sample=[radar,semseg] \
+        objective=semseg \
+        model/decoder=semseg
+    ```
 
-    === "Ego-Motion"
-        ```sh
-        uv run train.py meta.name=vel meta.version=small size=small \
-            +base=occ3d_to_vel \
-            sensors=[radar,pose]
-            transforms@transforms.sample=[radar,vel] \
-            objective=vel \
-            model/decoder=vel
-        ```
+=== "Ego-Motion"
+    ```sh
+    uv run train.py meta.name=vel meta.version=small size=small \
+        +base=occ3d_to_vel \
+        sensors=[radar,pose] \
+        transforms@transforms.sample=[radar,vel] \
+        objective=vel \
+        model/decoder=vel
+    ```
 
 !!! tip
 
@@ -119,6 +119,13 @@ The GRT template includes reference training scripts which can be used for high
     --8<-- "grt/train_minimal.py"
     ```
 
+!!! info "Compile the Model"
+
+    You can invoke the pytorch [JIT Compiler](https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html) by setting `meta.compile=True`. Since the pytorch compiler is kind of janky and causes issues with type checkers, you will also need to set
+    ```sh
+    export JAXTYPING_DISABLE=1
+    ```
+
 ## Evaluation Script
 
 ::: evaluate.evaluate

docs/index.md

@@ -55,31 +55,33 @@ The NRDK is intended to be used as a library; we recommend installing it as a su
 
     ```sh
     git submodule add git@github.com:RadarML/nrdk.git
-    uv pip install -e ./nrdk[roverd,grt]
+    uv pip install -e ./nrdk[roverd]
     ```
 
-=== "Development"
+=== "Via `uv` + `pyproject.toml`"
 
-    ```sh
-    git clone git@github.com:RadarML/nrdk.git
-    cd nrdk
-    uv sync --all-extras
-    uv run pre-commit install
+    ```toml
+    [project]
+    dependencies = ["nrdk[roverd]"]
+
+    [tool.uv.extensions]
+    nrdk = { path = "./nrdk", editable = true }
     ```
 
-!!! tip
+=== "For Development"
 
-    If authenticating with a github token (i.e., `https`), translate all `ssh://git@github.com` authentication to `https://github.com` with
     ```sh
-    git config --global url."https://github.com/".insteadOf ssh://git@github.com/
+    git clone git@github.com:RadarML/nrdk.git
+    cd nrdk
+    uv sync --all-extras --frozen
+    uv run pre-commit install
     ```
 
 ??? info "Extras"
 
     The NRDK library also includes the following extras:
 
     - `nrdk[roverd]`: support for loading and processing data using the [roverd](https://radarml.github.io/red-rover/roverd/) format (i.e., as collected by the [`red-rover`](https://radarml.github.io/red-rover/) system).
-    - `nrdk[grt]`: extra dependencies for our GRT reference implementation, which uses a hydra-based configuration system.
     - `nrdk[docs]`: a mkdocs + mkdocs-material + mdocstrings-python documentation stack.
     - `nrdk[dev]`: linters, testing, pre-commit hooks, etc.
 
@@ -123,6 +125,12 @@ uv run pre-commit run --all-files
 
     abstract interface for composable dataloaders and preprocessing pipelines
 
+- :fontawesome-solid-hexagon-nodes: [`grt`](https://wiselabcmu.github.io/grt/)
+
+    --- 
+
+    *our latest work, GRT: Towards Foundational Models for Single-Chip Radar*
+
 - :dart: [`dart`](https://wiselabcmu.github.io/dart/)
 
     ---

grt/config/default.yaml

@@ -13,13 +13,18 @@ defaults:
 hydra:
   run:
     dir: ${meta.results}/${meta.name}/${meta.version}
+  callbacks:
+    prevent_overwrite:
+      _target_: nrdk.config.PreventHydraOverwrite
+
 
 meta:
-  dataset: data/
-  results: results/
+  dataset: data
+  results: results
   name: null
   version: null
   resume: null
+  compile: false
 
 size: {}
 
@@ -57,7 +62,6 @@ datamodule:
 
 lightningmodule:
   _target_: nrdk.framework.NRDKLightningModule
-  compile: false
   vis_interval: 1000
   vis_samples: 16
   objective: ${objective}

grt/config/model/decoder/lidar3d.yaml

@@ -12,7 +12,7 @@ occ3d:
     norm_first: true
     bias: true
   d_model: ${size.d_model}
-  num_layers: 4
+  num_layers: ${size.dec_layers}
   shape: [1, 64, 128, 64]
   scale: [16.0, 16.0, 16.0, 16.0]
   w_min: 0.2

grt/config/model/decoder/semseg.yaml

@@ -12,7 +12,7 @@ semseg:
     norm_first: true
     bias: true
   d_model: ${size.d_model}
-  num_layers: 4
+  num_layers: ${size.dec_layers}
   shape: [1, 160, 160, 1]
   scale: [16.0, 16.0, 10.666, 16.0]
   w_min: 0.2

grt/config/objective/lidar2d.yaml

@@ -6,7 +6,7 @@ occupancy2d:
   objective:
     _target_: nrdk.objectives.Occupancy2D
     range_weighted: True
-    positive_weight: 64.0
+    positive_weight: 1.0
     bce_weight: 0.9
     vis_config:
       cols: 8

grt/config/objective/lidar3d.yaml

@@ -6,7 +6,7 @@ occupancy3d:
   objective:
     _target_: nrdk.objectives.Occupancy3D
     range_weighted: True
-    positive_weight: 64.0
+    positive_weight: 16.0
     mode: spherical
     vis_config:
       cols: 8