Misc refactoring

Author: thetianshuhuang

Makefile

@@ -37,7 +37,7 @@ $(_SLICES): $(_MAP)
 $(_RAD): $(_TGT)
 	$(DART) evaluate -p results/$(TARGET) -a -b $(BATCH)
 # SSIM calculation
-$(_SSIM): $(_TGT)
+$(_SSIM): $(_RAD)
 	$(DART) ssim -p results/$(TARGET)
 # Camera evaluation
 $(_CAM): $(_TGT)
@@ -54,7 +54,7 @@ train: results/$(TARGET)
 map: $(_MAP)
 slices: $(_SLICES)
 radar: $(_RAD)
-ssim: results/$(TARGET)/ssim.npz
+ssim: $(_SSIM)
 camera: $(_CAM)
 video: $(_VIDEO)
 

dart/dataset.py

@@ -46,13 +46,6 @@ def load_arrays(file: str, keys: Optional[list[str]] = None) -> Any:
             "Unknown file type: {} (expected .npz, .h5, or .mat)".format(file))
 
 
-def gt_map(file: str) -> GroundTruth:
-    """Load ground truth reflectance map."""
-    data = load_arrays(file)
-    return GroundTruth.from_occupancy(
-        jnp.array(data['grid']), data["lower"], data["upper"])
-
-
 def trajectory(
     traj: str, subset: Optional[Integer[types.ArrayLike, "nval"]] = None
 ) -> types.Dataset:
@@ -113,8 +106,8 @@ def __doppler_decimation(
 
 
 def doppler_columns(
-    sensor: VirtualRadar, path: str,
-    pval: float = 0., iid_val: bool = False, doppler_decimation: int = 0,
+    path: str, pval: float = 0., iid_val: bool = False,
+    doppler_decimation: int = 0,
     key: types.PRNGSeed = 42
 ) -> tuple[types.Dataset, Optional[types.Dataset], dict[str, PyTree]]:
     """Load dataset trajectory and images.
@@ -129,55 +122,50 @@ def doppler_columns(
 
     Parameters
     ----------
-    sensor: Sensor profile for this dataset.
     path: Path to file containing data.
-    norm: Normalization factor.
     pval: Proportion of dataset to hold as a validation set. If `pval=0`,
         no validation dataset is returned.
     iid_val: If True, then shuffles the dataset before training so that the
         validation split is drawn randomly from the dataset instead of just
         from the end.
-    min_speed: Minimum speed for usable samples. Images with lower
-        velocities are rejected.
-    repeat: Repeat dataset within each epoch to reduce overhead.
-    threshold: Mask out values less than the provided threshold (set to 0).
-    doppler_decimation: Simulate a lower doppler resolution by setting each
-        block of consecutive doppler columns to their average.
     key: Random key to shuffle dataset frames. Does not shuffle columns.
 
     Returns
     -------
     train: Train dataset.
     val: Val dataset.
-    validx: Indices of original images corresponding to the validation set.
+    meta: Metadata (exact split indices).
     """
     file = h5py.File(path)
 
     pose = types.RadarPose.from_h5file(file)
-    idx = np.array(file["idx"])
     rad = np.array(file["rad"], dtype=np.float16)
     weight = np.array(file["weight"], dtype=np.float32)
     doppler = np.array(file["doppler"], dtype=np.float32)
+    idx = np.arange(rad.shape[0])
 
     meta = types.TrainingColumn(pose=pose, weight=weight, doppler=doppler)
-    data = (meta, rad)
+    data = (meta, rad), idx
 
     print("Loaded dataset : {} valid columns".format(rad.shape))
 
     if iid_val:
         data = utils.shuffle(data, key=key)
 
-    meta, rad = data
-    nval = 0 if pval <= 0 else int(utils.get_size(data) * pval)
-    train, val = utils.split((meta, rad), nval=nval)
+    nval = 0 if pval <= 0 else int(rad.shape[0] * pval)
+    (train, itrain), _val = utils.split(data, nval=nval)
 
     if not iid_val:
         train = utils.shuffle(train, key=key)
 
     print("Train split    : {} columns".format(train[1].shape))
     train = types.Dataset.from_tensor_slices(train)
-    if val is not None:
+    if _val is not None:
+        val, ival = _val
         print("Test split     : {} columns".format(val[1].shape))
         val = types.Dataset.from_tensor_slices(val)
+    else:
+        val = None
+        ival = np.zeros(0, dtype=bool)
 
-    return train, val, {"val": idx[-nval:]}
+    return train, val, {"train": itrain, "val": ival}

dart/fields/ngp.py

@@ -13,8 +13,8 @@
 import jax
 import haiku as hk
 
-from jaxtyping import Float32, Integer, Array, Float
-from beartype.typing import Optional, Callable, Union
+from jaxtyping import Float32, Integer, Array
+from beartype.typing import Optional, Callable
 
 from dart import types
 from ._spatial import interpolate, spherical_harmonics
@@ -104,11 +104,15 @@ def hash_table(c):
 
     def __call__(
         self, x: Float32[Array, "3"], dx: Optional[Float32[Array, "3"]] = None,
-        **kwargs
+        alpha_clip: Optional[types.FloatLike] = None, **kwargs
     ) -> tuple[Float32[Array, ""], Float32[Array, ""]]:
         """Index into learned reflectance map."""
         table_out = self.lookup(x)
         sigma, alpha = self.head(table_out.reshape(-1))
+
+        if alpha_clip is not None:
+            alpha = jnp.where(sigma > alpha_clip, alpha, 0)
+
         return sigma, clip(alpha) * self.alpha_scale
 
     @classmethod
@@ -239,7 +243,7 @@ def __init__(self, harmonics: int = 16, **kwargs) -> None:
 
     def __call__(
         self, x: Float32[Array, "3"], dx: Optional[Float32[Array, "3"]] = None,
-        **kwargs
+        alpha_clip: Optional[types.FloatLike] = None, **kwargs
     ) -> tuple[Float32[Array, ""], Float32[Array, ""]]:
         """Index into learned reflectance map."""
         table_out = self.lookup(x)
@@ -249,6 +253,10 @@ def __call__(
 
         mlp_out = self.head(jnp.concatenate([table_out.reshape(-1), sh]))
         sigma, alpha = mlp_out
+
+        if alpha_clip is not None:
+            alpha = jnp.where(sigma > alpha_clip, alpha, 0)
+
         return sigma, clip(alpha) * self.alpha_scale
 
     @staticmethod

dart/pose.py

@@ -55,7 +55,7 @@ def make_pose(
 
 
 def sensor_to_world(
-    r: Float32[Array, ""], t: Float32[Array, "3 k"],
+    r: Float32[types.ArrayLike, ""], t: Float32[types.ArrayLike, "3 k"],
     pose: Union[types.CameraPose, types.RadarPose]
 ) -> Float32[Array, "3 k"]:
     """Project points to world-space.
@@ -74,7 +74,8 @@ def sensor_to_world(
 
 
 def project_angle(
-    d: Float32[Array, ""], psi: Float32[Array, "n"], pose: types.RadarPose
+    d: Float32[types.ArrayLike, ""], psi: Float32[types.ArrayLike, "n"],
+    pose: types.RadarPose
 ) -> Float32[Array, "3 n"]:
     """Project angles to intersection circle on a unit sphere.
 

dart/script.py

@@ -24,7 +24,7 @@ def script_train(cfg: dict) -> None:
     k1, k2, k3 = jax.random.split(root, 3)
 
     dart = DART.from_config(**cfg)
-    train, val, meta = doppler_columns(dart.sensor, key=k1, **cfg["dataset"])
+    train, val, meta = doppler_columns(key=k1, **cfg["dataset"])
     assert val is not None
     train = train.shuffle(cfg["shuffle_buffer"], reshuffle_each_iteration=True)
 

dart/sensor.py

@@ -58,7 +58,7 @@ def from_config(
     @staticmethod
     def get_psi_min(
         d: Float32[types.ArrayLike, ""], pose: types.RadarPose
-    ) -> Float32[types.ArrayLike, ""]:
+    ) -> Float32[Array, ""]:
         """Get psi value representing visible region of integration circle.
 
         Visible psi angles fall in the range of (-psi_min, psi_min). These

dart/utils/jaxcolors.py

@@ -7,13 +7,23 @@
 
 
 def hsv_to_rgb(
-    hsv: Float[types.ArrayLike, "... 3"]
+    hsv: Float[types.ArrayLike, "... 3"], _np=jnp
 ) -> Float[Array, "... 3"]:
     """Convert hsv values to rgb.
 
     Copied here, and modified for vectorization:
     https://matplotlib.org/3.1.1/_modules/matplotlib/colors.html#hsv_to_rgb
     and converted to jax.
+
+    Parameters
+    ----------
+    hsv: HSV colors.
+    _np: numpy-like backend to use.
+
+    Returns
+    -------
+    RGB colors `float: (0, 1)`, using the array format corresponding to the
+    provided backend.
     """
     in_shape = hsv.shape
     h = hsv[..., 0]
@@ -36,10 +46,23 @@ def hsv_to_rgb(
 
 
 def colormap(
-    colors: Num[types.ArrayLike, "n 3"],
-    data: Float[types.ArrayLike, "..."]
+    colors: Num[types.ArrayLike, "n d"],
+    data: Float[types.ArrayLike, "..."],
+    _np=jnp
 ) -> Num[Array, "... 3"]:
-    """Apply a discrete colormap."""
+    """Apply a discrete colormap.
+    
+    Parameters
+    ----------
+    colors: list of discrete colors to apply (e.g. 0-255 RGB values). Can be
+        an arbitrary number of channels, not just RGB.
+    data: input data to index (`0 < data < 1`).
+    _np: numpy-like backend to use.
+
+    Returns
+    -------
+    An array with the same shape as `data`, with an extra dimension appended.
+    """
     fidx = data * (colors.shape[0] - 1)
-    left = jnp.clip(jnp.floor(fidx).astype(int), 0, colors.shape[0] - 1)
-    return jnp.take(colors, left, axis=0)
+    left = _np.clip(_np.floor(fidx).astype(int), 0, colors.shape[0] - 1)
+    return _np.take(colors, left, axis=0)

dart/utils/misc.py

@@ -10,14 +10,17 @@
 from jax import numpy as jnp
 import numpy as np
 
-
 from jaxtyping import PyTree
 from beartype.typing import TypeVar, Optional, Union
 from dart import types
 
 
 def tf_to_jax(batch: PyTree) -> PyTree:
-    """Convert tensorflow array to jax array without copying."""
+    """Convert tensorflow array to jax array without copying.
+    
+    NOTE: going through dlpack is much slower, so it seems jax/tf have some
+    kind of interop going on under the hood already.
+    """
     return jax.tree_util.tree_map(jnp.array, batch)
 
 

dart/utils/ssim.py

@@ -4,7 +4,7 @@
 from jax import numpy as jnp
 import jax.scipy as jsp
 
-from jaxtyping import Float, Float32, Array
+from jaxtyping import Float, Array
 
 
 def ssim(

datasets.md

@@ -0,0 +1,9 @@
+# Hyperparameter Tuning Notes
+
+`alpha_clip` (`--clip`): 0.05 seems to work pretty well.
+- If the velocity is very noisy / bad, high `alpha_clip` can cause problems.
+- Low `alpha_clip` causes phantom "trails" behind the trajectory of high reflectance + low transmittance, which probably indicates overfitting.
+
+`epochs` (`--epochs`): 3 seems to be enough.
+- With an appropriate `alpha_clip`, the val loss doesn't seem to move by much after the first epoch.
+- Too many epochs can cause overfitting which is qualitatively apparent on a visual inspection (the loss doesn't really tell the whole story). The reflectance map starts to get a lot of holes, and the transmittance map is even less continuous.

docs/hyperparams.md

@@ -12,8 +12,7 @@
     subparsers = parser.add_subparsers()
     for name, command in commands.items():
         p = subparsers.add_parser(
-            name, help=command._desc.split('\n')[0],
-            description=command._desc,
+            name, help=command.__doc__, description=command.__doc__,
             formatter_class=RawTextHelpFormatter)
         command._parse(p)
         p.set_defaults(_func=command._main)