Merge pull request #18 from RadarML/dev/modules

GRT Model Fixes & Model Features

Author: thetianshuhuang

docs/grt/index.md

@@ -14,6 +14,49 @@ The GRT reference implementation uses a [hydra](https://hydra.cc/docs/intro/) +
 
     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:
+
+    === "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
+        ```
+
+    === "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
+        ```
+
+!!! tip
+
+    If you're not running in a "managed" environment (e.g., Slurm, LSF, AzureML), [`nq`](https://github.com/leahneukirchen/nq) is a lightweight way to run jobs in a queue. Just `sudo apt-get install -y nq`, and run with `nq uv run train.py ...`.
+
+
 ## Quick Start
 
 1. Create a new repository, and copy the contents of the `grt/` directory:

grt/config/base/occ3d_to_occ2d.yaml

@@ -1,4 +1,4 @@
-path: results/example/base/weights.pth
+path: results/base/small/weights.pth
 rename:
 - decoder.occ3d.unpatch: null
 - decoder.occ3d: decoder.occ2d

grt/config/base/occ3d_to_semseg.yaml

@@ -1,4 +1,4 @@
-path: results/example/base/weights.pth
+path: results/base/small/weights.pth
 rename:
 - decoder.occ3d.unpatch: null
 - decoder.occ3d: decoder.semseg

grt/config/base/occ3d_to_vel.yaml

@@ -1,3 +1,3 @@
-path: results/example/base/weights.pth
+path: results/base/small/weights.pth
 rename:
 - decoder: null

grt/config/default.yaml

@@ -2,12 +2,17 @@ defaults:
   - _self_
   - size: small
   - aug@transforms.sample.augmentations: full
-  - sensors@datamodule.dataset.sensors: ["radar", "lidar"]
+  - sensors: ["radar", "lidar"]
   - transforms@transforms.sample: ["radar", "lidar3d"]
   - traces@datamodule.traces: ["bike", "indoor", "outdoor"]
   - model: default
   - trainer: default
   - objective: lidar3d
+  - optimizer: default
+
+hydra:
+  run:
+    dir: ${meta.results}/${meta.name}/${meta.version}
 
 meta:
   dataset: data/
@@ -18,9 +23,11 @@ meta:
 
 size: {}
 
-hydra:
-  run:
-    dir: ${meta.results}/${meta.name}/${meta.version}
+model: {}
+objective: {}
+aug: {}
+optimizer: {}
+sensors: {}
 
 datamodule:
   _target_: nrdk.roverd.datamodule
@@ -31,15 +38,15 @@ datamodule:
       _target_: abstract_dataloader.generic.Next
       margin: [1.0, 1.0]
       reference: "radar"
-    sensors: {}  # sensors/*.yaml
+    sensors: ${sensors}
   traces:
     train:
       _target_: nrdk.config.expand
       path: ${meta.dataset}
     test:
       _target_: nrdk.config.expand
       path: ${meta.dataset}
-  batch_size: 32
+  batch_size: 64
   samples: 8
   num_workers: 16
   prefetch_factor: 1
@@ -48,31 +55,22 @@ datamodule:
   ptrain: 0.8
   pval: 0.2
 
-model: {}
-
-objective: {}
-
 lightningmodule:
   _target_: nrdk.framework.NRDKLightningModule
   compile: false
   vis_interval: 1000
   vis_samples: 16
   objective: ${objective}
   model: ${model}
-  optimizer:
-    _target_: torch.optim.AdamW
-    _partial_: true
-    lr: 1e-4
-    weight_decay: 1e-2
-    betas: [0.9, 0.999]
-    eps: 1e-8
+  optimizer: ${optimizer}
 
 transforms:
   _target_: abstract_dataloader.abstract.Pipeline
   sample:
     _target_: abstract_dataloader.ext.graph.Transform
     keep_all: false
     outputs: {}  # sensors/*.yaml
+    augmentations: ${aug}
   collate:
     _target_: abstract_dataloader.ext.torch.Collate
     mode: "concat"

grt/config/model/decoder/semseg.yaml

@@ -17,4 +17,4 @@ semseg:
   scale: [16.0, 16.0, 10.666, 16.0]
   w_min: 0.2
   patch: [1, 5, 5, 1]
-  out_dim: 8
+  out_dim: 8

grt/config/objective/lidar2d.yaml

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

grt/config/optimizer/default.yaml

@@ -0,0 +1,6 @@
+_target_: torch.optim.AdamW
+_partial_: true
+lr: 1e-4
+weight_decay: 1e-2
+betas: [0.9, 0.999]
+eps: 1e-8

grt/train.py

@@ -62,6 +62,14 @@ def _inst(path, *args, **kwargs):
         return hydra.utils.instantiate(
             cfg[path], _convert_="all", *args, **kwargs)
 
+    n_gpus = torch.cuda.device_count()
+    if "batch_size" in cfg["datamodule"] and n_gpus > 1:
+        batch_new = cfg["datamodule"]["batch_size"] // n_gpus
+        logger.info(
+            f"Auto-scaling batch size by n_gpus={n_gpus}: "
+            f"{cfg["datamodule"]["batch_size"]} -> {batch_new}")
+        cfg["datamodule"]["batch_size"] = batch_new
+
     transforms = _inst("transforms")
     datamodule = _inst("datamodule", transforms=transforms)
     lightningmodule = _inst("lightningmodule", transforms=transforms)