nn wrapper; v4; test

Author: misko

spf/dataset/spf_dataset.py

@@ -421,6 +421,7 @@ def __setitem__(self, key, value):
     "ground_truth_theta",
     "ground_truth_phi",
     "craft_ground_truth_theta",
+    "absolute_theta",
     "y_rad",
     "y_phi",
     "craft_y_rad",
@@ -429,7 +430,7 @@ def __setitem__(self, key, value):
 
 segmentation_based_keys = [
     "weighted_beamformer",
-    "all_windows_stats",
+    # "all_windows_stats",
     "weighted_windows_stats",
     "downsampled_segmentation_mask",
     "simple_segmentations",
@@ -439,8 +440,8 @@ def __setitem__(self, key, value):
 v5_raw_keys = v5rx_f64_keys + v5rx_2xf64_keys + ["signal_matrix"]
 
 
-def data_single_radio_to_raw(d):
-    return {k: d[k] for k in v5_raw_keys}
+def data_single_radio_to_raw(d, ds):
+    return {k: d[k] for k in list(set(v5_raw_keys) - set(ds.skip_fields))}
 
 
 class v5inferencedataset(Dataset):
@@ -470,13 +471,16 @@ def __init__(
         skip_segmentation: bool = True,
         vehicle_type: str = "",
         max_in_memory: int = 10,
+        realtime: bool = True,
     ):
         # Store configuration parameters
         self.yaml_fn = yaml_fn
         self.n_parallel = n_parallel
         self.nthetas = nthetas  # Number of angles to discretize space for beamforming
         self.target_ntheta = self.nthetas if target_ntheta is None else target_ntheta
 
+        self.realtime = realtime
+
         self.max_in_memory = max_in_memory
         self.min_idx = 0
         self.condition = multiprocessing.Condition()
@@ -589,6 +593,9 @@ def __init__(
             self.empirical_data_fn = None
             self.empirical_data = None
 
+    def __len__(self):
+        return self.serving_idx
+
     def __iter__(self):
         self.serving_idx = 0
         return self
@@ -606,9 +613,11 @@ def __next__(self):
     def __getitem__(self, idx, timeout=10.0):
         start_time = time.time()
         with self.condition:
+            print("waitinf to get get", idx, time.time() - start_time)
             while idx not in self.store or self.store[idx]["count"] != 2:
                 self.condition.wait(0.01)
                 if (time.time() - start_time) > timeout:
+                    print("ret waitinf to get get", idx, time.time() - start_time)
                     return None
             return self.store[idx]["data"]
 
@@ -637,6 +646,7 @@ def render_session(self, idx, ridx, data):
         data["receiver_idx"] = torch.tensor([[ridx]], dtype=torch.int)
 
         data["ground_truth_theta"] = torch.tensor([torch.inf])  # unknown
+        data["absolute_theta"] = torch.tensor([torch.inf])
         data["y_rad"] = data["ground_truth_theta"]  # torch.inf
 
         data["ground_truth_phi"] = torch.tensor([torch.inf])  # unkown
@@ -774,6 +784,7 @@ def __init__(
         skip_detrend: bool = False,
         vehicle_type: str = "",
         v4: bool = False,
+        realtime: bool = False,
     ):
         logging.debug(f"loading... {prefix}")
         # Store configuration parameters
@@ -786,6 +797,8 @@ def __init__(
         self.valid_entries = None
         self.temp_file = temp_file
 
+        self.realtime = realtime
+
         # Segmentation parameters control how raw signal is processed into windows
         # and how phase difference is computed between antenna elements
         self.segmentation_version = (
@@ -1286,6 +1299,7 @@ def render_session(self, receiver_idx, session_idx, double_flip=False):
         data["ground_truth_theta"] = self.ground_truth_thetas[receiver_idx][
             snapshot_idxs
         ]
+        data["absolute_theta"] = self.absolute_thetas[receiver_idx][snapshot_idxs]
         data["ground_truth_phi"] = self.ground_truth_phis[receiver_idx][snapshot_idxs]
         data["craft_ground_truth_theta"] = self.craft_ground_truth_thetas[snapshot_idxs]
         data["vehicle_type"] = torch.tensor(

spf/dataset/spf_nn_dataset_wrapper.py

@@ -0,0 +1,146 @@
+from functools import cache
+
+import torch
+from torch.utils.data import Dataset
+
+from spf.dataset.spf_dataset import v5_collate_keys_fast
+from spf.filters.particle_dual_radio_nn_filter import (
+    cached_model_inference_to_absolute_north,
+)
+from spf.model_training_and_inference.models.single_point_networks_inference import (
+    get_nn_inference_on_ds_and_cache,
+    load_model_and_config_from_config_fn_and_checkpoint,
+)
+from spf.rf import rotate_dist, torch_pi_norm_pi
+from spf.scripts.train_utils import global_config_to_keys_used
+
+
+# attach nn inference attributes to dataset entries
+class v5spfdataset_nn_wrapper(Dataset):
+    def __init__(
+        self,
+        ds,
+        checkpoint_config_fn,
+        checkpoint_fn,
+        inference_cache,
+        device="cpu",
+        v4=None,
+        absolute=False,
+    ):
+        self.ds = ds
+        assert self.ds.paired
+        self.checkpoint_config_fn = checkpoint_config_fn
+        self.checkpoint_fn = checkpoint_fn
+        self.inference_cache = inference_cache
+        self.absolute = absolute
+
+        if v4 is None:
+            v4 = self.ds.v4
+
+        if not ds.realtime:
+            self.cached_model_inference = {
+                k: torch.as_tensor(v)
+                for k, v in get_nn_inference_on_ds_and_cache(
+                    ds_fn=ds.zarr_fn,
+                    config_fn=self.checkpoint_config_fn,
+                    checkpoint_fn=self.checkpoint_fn,
+                    device=device,
+                    inference_cache=inference_cache,
+                    batch_size=64,
+                    workers=0,
+                    precompute_cache=ds.precompute_cache,
+                    crash_if_not_cached=False,
+                    segmentation_version=ds.segmentation_version,
+                    v4=v4,
+                ).items()
+            }
+            if self.absolute:
+                self.cached_model_inference = {
+                    k: cached_model_inference_to_absolute_north(ds, v)
+                    for k, v in self.cached_model_inference.items()
+                }
+        else:
+            self.model, self.model_config = (
+                load_model_and_config_from_config_fn_and_checkpoint(
+                    self.checkpoint_config_fn, self.checkpoint_fn, device=device
+                )
+            )
+            self.model.eval()
+            self.keys_to_get = global_config_to_keys_used(
+                global_config=self.model_config["global"]
+            )
+
+    def to_absolute_north(self, sample):
+        for ridx in range(2):
+            ntheta = sample[ridx]["paired"].shape[-1]
+            paired_nn_inference = sample[ridx]["paired"].reshape(-1, ntheta)
+            paired_nn_inference_rotated = rotate_dist(
+                paired_nn_inference,
+                rotations=torch_pi_norm_pi(
+                    sample[ridx]["rx_heading_in_pis"][:, None] * torch.pi
+                ),
+            ).reshape(paired_nn_inference.shape)
+            sample[ridx]["paired"] = paired_nn_inference_rotated
+        return sample
+
+    @cache
+    def get_inference_for_idx(self, idx):
+        if not self.ds.realtime:
+            return [
+                {k: v[idx][ridx] for k, v in self.cached_model_inference.items()}
+                for ridx in range(2)
+            ]
+        return self.get_and_annotate_entry_at_idx(idx)
+
+    @cache
+    def get_and_annotate_entry_at_idx(self, idx):
+        sample = self.ds[idx]
+        if not self.ds.realtime:
+            for ridx in range(2):
+                sample[ridx].update(
+                    {k: v[idx][ridx] for k, v in self.cached_model_inference.items()}
+                )
+            return sample
+        else:
+            single_example = v5_collate_keys_fast(self.keys_to_get, [sample]).to(
+                self.model_config["optim"]["device"]
+            )
+            with torch.no_grad():
+                nn_inference = self.model(single_example)
+            for ridx in range(2):
+                sample[ridx].update({k: v[ridx] for k, v in nn_inference.items()})
+            if self.absolute:
+                sample = self.to_absolute_north(sample)
+            return sample
+
+    def __iter__(self):
+        self.serving_idx = 0
+        return self
+
+    def __next__(self):
+        sample = self.get_and_annotate_entry_at_idx(self.serving_idx)
+        self.serving_idx += 1
+        return sample
+
+    @cache
+    def __getitem__(self, idx):
+        return self.get_and_annotate_entry_at_idx(idx)
+
+    def __len__(self):
+        return len(self.ds)
+
+    @property
+    def mean_phase(self):
+        return self.ds.mean_phase
+
+    @property
+    def ground_truth_phis(self):
+        return self.ds.ground_truth_phis
+
+    @property
+    def craft_ground_truth_thetas(self):
+        return self.ds.craft_ground_truth_thetas
+
+    @property
+    def absolute_thetas(self):
+        return self.ds.absolute_thetas

spf/filters/particle_dual_radio_nn_filter.py

@@ -1,22 +1,12 @@
 import torch
 
-from spf.dataset.spf_dataset import v5_collate_keys_fast
 from spf.filters.filters import (
     ParticleFilter,
     add_noise,
     dual_radio_mse_theta_metrics,
     theta_phi_to_bins,
 )
-from spf.model_training_and_inference.models.single_point_networks_inference import (
-    convert_datasets_config_to_inference,
-    get_inference_on_ds,
-    load_model_and_config_from_config_fn_and_checkpoint,
-)
 from spf.rf import rotate_dist, torch_pi_norm_pi
-from spf.scripts.train_single_point import (
-    global_config_to_keys_used,
-    load_config_from_fn,
-)
 
 
 def cached_model_inference_to_absolute_north(ds, cached_model_inference):
@@ -40,80 +30,15 @@ def cached_model_inference_to_absolute_north(ds, cached_model_inference):
 class PFSingleThetaDualRadioNN(ParticleFilter):
     def __init__(
         self,
-        ds,
-        checkpoint_fn,
-        config_fn,
-        inference_cache=None,
-        device="cpu",
-        absolute=False,
+        nn_ds,
     ):
-        self.ds = ds
-        self.absolute = absolute
+        self.ds = nn_ds
+        self.absolute = nn_ds.absolute
         self.generator = torch.Generator()
         self.generator.manual_seed(0)
 
-        # checkpoint_config = load_config_from_fn(config_fn)
-        # assert (
-        #     self.ds.empirical_data_fn
-        #     == checkpoint_config["datasets"]["empirical_data_fn"]
-        # )
-
-        if not self.ds.temp_file:
-            # cache model results
-            self.cached_model_inference = torch.as_tensor(
-                get_inference_on_ds(
-                    ds_fn=ds.zarr_fn,
-                    config_fn=config_fn,
-                    checkpoint_fn=checkpoint_fn,
-                    device=device,
-                    inference_cache=inference_cache,
-                    batch_size=64,
-                    workers=0,
-                    precompute_cache=ds.precompute_cache,
-                    crash_if_not_cached=False,
-                    segmentation_version=ds.segmentation_version,
-                )["paired"]
-            )
-            if self.absolute:
-                self.cached_model_inference = cached_model_inference_to_absolute_north(
-                    ds, self.cached_model_inference
-                )
-        else:
-            # load the model and such
-            self.model, self.model_config = (
-                load_model_and_config_from_config_fn_and_checkpoint(
-                    config_fn=config_fn, checkpoint_fn=checkpoint_fn, device=device
-                )
-            )
-            self.model.eval()
-
-            self.model_datasets_config = convert_datasets_config_to_inference(
-                self.model_config["datasets"],
-                ds_fn=ds.zarr_fn,
-                precompute_cache=self.ds.precompute_cache,
-            )
-
-            self.model_optim_config = {"device": device, "dtype": torch.float32}
-
-            self.model_keys_to_get = global_config_to_keys_used(
-                global_config=self.model_config["global"]
-            )
-            assert not self.absolute  # this needs to be implemented
-
-    def model_inference_at_observation_idx(self, idx):
-        if not self.ds.temp_file:
-            return self.cached_model_inference[idx]
-
-        z = v5_collate_keys_fast(self.model_keys_to_get, [self.ds[idx]]).to(
-            self.model_optim_config["device"]
-        )
-        with torch.no_grad():
-            return self.model(z)["paired"].cpu()
-
     def observation(self, idx):
-        # even though the model outputs one paired dist for each reciever
-        # they should be identical
-        return self.model_inference_at_observation_idx(idx)[0, 0]
+        return self.ds.get_inference_for_idx(idx)[0]["paired"][0]
 
     def fix_particles(self):
         self.particles[:, 0] = torch_pi_norm_pi(self.particles[:, 0])
@@ -125,7 +50,6 @@ def predict(self, our_state, dt, noise_std):
         add_noise(self.particles, noise_std=noise_std, generator=self.generator)
 
     def update(self, z):
-        #
         # z is not the raw observation, but the processed model output
         theta_bin = theta_phi_to_bins(self.particles[:, 0], nbins=z.shape[0])
         prob_theta_given_observation = torch.take(z, theta_bin)

spf/filters/particle_single_radio_nn_filter.py

@@ -10,7 +10,7 @@
 )
 from spf.model_training_and_inference.models.single_point_networks_inference import (
     convert_datasets_config_to_inference,
-    get_inference_on_ds,
+    get_nn_inference_on_ds_and_cache,
     load_model_and_config_from_config_fn_and_checkpoint,
 )
 from spf.scripts.train_single_point import (
@@ -41,7 +41,7 @@ def __init__(
         if not self.ds.temp_file:
             # cache model results
             self.cached_model_inference = torch.as_tensor(
-                get_inference_on_ds(
+                get_nn_inference_on_ds_and_cache(
                     ds_fn=ds.zarr_fn,
                     config_fn=config_fn,
                     checkpoint_fn=checkpoint_fn,