downsample lerf samples

nerfstudio/configs/method_configs.py

@@ -96,7 +96,7 @@ method_configs["lerf"] = TrainerConfig(
     mixed_precision=True,
     pipeline=LERFPipelineConfig(
         datamanager=LERFDataManagerConfig(
-            dataparser=NerfstudioDataParserConfig(),
+            dataparser=NerfstudioDataParserConfig(train_split_fraction=0.99),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,
             camera_optimizer=CameraOptimizerConfig(
@@ -108,15 +108,15 @@ method_configs["lerf"] = TrainerConfig(
     optimizers={
         "proposal_networks": {
             "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": None,
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000),
         },
         "fields": {
             "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": None,
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000),
         },
         "lerf": {
             "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15),
-            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=5000),
+            "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=7000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

nerfstudio/models/lerf.py

@@ -5,13 +5,13 @@ import numpy as np
 import open_clip
 import torch
 from torch.nn import Parameter
-from typing_extensions import Literal
 
 from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.lerf_field import LERFField
+from nerfstudio.model_components.ray_samplers import PDFSampler
 from nerfstudio.model_components.renderers import (
     CLIPRenderer,
     DepthRenderer,
@@ -28,6 +28,7 @@ class LERFModelConfig(NerfactoModelConfig):
     max_scale: float = 1.5
     """maximum scale used to compute relevancy with"""
     specify_scale: bool = False
+    num_lerf_samples: int = 12
 
 
 class LERFModel(NerfactoModel):
@@ -41,6 +42,7 @@ class LERFModel(NerfactoModel):
 
         self.network = self.kwargs["network"]
         self.lerf_field = LERFField(clip_n_dims=self.network.clip_n_dims)
+        # self.lerf_sampler = PDFSampler(num_samples=self.config.num_lerf_samples, histogram_padding=0.00)
 
     def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, preset_scales=None):
         # TODO smoothen this out
@@ -78,15 +80,28 @@ class LERFModel(NerfactoModel):
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         ray_samples_list.append(ray_samples)
 
+        nerfacto_field_outputs, outputs, weights = self._get_outputs_nerfacto(ray_samples)
+        lerf_weights, best_ids = torch.topk(weights, self.config.num_lerf_samples, dim=-2, sorted=False)
+
+        def gather_fn(tens):
+            return torch.gather(tens, -2, best_ids.expand(*best_ids.shape[:-1], tens.shape[-1]))
+
+        dataclass_fn = lambda dc: dc._apply_fn_to_fields(gather_fn, dataclass_fn)
+        lerf_samples = ray_samples._apply_fn_to_fields(gather_fn, dataclass_fn)
+
+        # lerf_samples = self.lerf_sampler(ray_bundle, ray_samples, weights)
+        # with torch.no_grad():
+        #     dens, _ = self.field.get_density(lerf_samples)
+        #     lerf_weights = lerf_samples.get_weights(dens)
+
         if "clip_scales" in ray_bundle.metadata:
             clip_scales = ray_bundle.metadata["clip_scales"]
             clip_scales = clip_scales[..., None]
-            dist = ray_samples.spacing_to_euclidean_fn(ray_samples.spacing_starts.squeeze(-1)).unsqueeze(-1)
+            dist = lerf_samples.spacing_to_euclidean_fn(lerf_samples.spacing_starts.squeeze(-1)).unsqueeze(-1)
             clip_scales = clip_scales * ray_bundle.metadata["width"] * (1 / ray_bundle.metadata["fx"]) * dist
         else:
-            clip_scales = torch.ones_like(ray_samples.spacing_starts, device=self.device)
+            clip_scales = torch.ones_like(lerf_samples.spacing_starts, device=self.device)
 
-        nerfacto_field_outputs, outputs, weights = self._get_outputs_nerfacto(ray_samples)
         weights_list.append(weights)
         if self.training:
             outputs["weights_list"] = weights_list
@@ -94,13 +109,18 @@ class LERFModel(NerfactoModel):
         for i in range(self.config.num_proposal_iterations):
             outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
 
-        lerf_field_outputs = self.lerf_field.get_outputs(ray_samples, clip_scales)
-        outputs["clip"] = self.renderer_clip(embeds=lerf_field_outputs[FieldHeadNames.CLIP], weights=weights.detach())
-        outputs["dino"] = self.renderer_mean(embeds=lerf_field_outputs[FieldHeadNames.DINO], weights=weights.detach())
+        lerf_field_outputs = self.lerf_field.get_outputs(lerf_samples, clip_scales)
+
+        outputs["clip"] = self.renderer_clip(
+            embeds=lerf_field_outputs[FieldHeadNames.CLIP], weights=lerf_weights.detach()
+        )
+        outputs["dino"] = self.renderer_mean(
+            embeds=lerf_field_outputs[FieldHeadNames.DINO], weights=lerf_weights.detach()
+        )
 
         if not "clip_scales" in ray_bundle.metadata:
             max_across, best_scales = self.get_max_across(
-                ray_samples, weights, lerf_field_outputs[FieldHeadNames.HASHGRID], clip_scales.shape
+                lerf_samples, lerf_weights, lerf_field_outputs[FieldHeadNames.HASHGRID], clip_scales.shape
             )
             multiphrase = max_across[0]
             # normalization for sanity check TODO(cmk) remove

nerfstudio/pipelines/lerf_pipeline.py

@@ -1,21 +1,27 @@
 import typing
-from typing import Literal, Type, List, Tuple
-
 from dataclasses import dataclass, field
-
-from nerfstudio.pipelines.base_pipeline import VanillaPipeline, VanillaPipelineConfig, Pipeline
-
-from nerfstudio.configs import base_config as cfg
-
-from nerfstudio.data.datamanagers.base_datamanager import DataManagerConfig, VanillaDataManagerConfig
-from nerfstudio.data.datamanagers.lerf_datamanager import LERFDataManager, LERFDataManagerConfig
-
-from nerfstudio.models.base_model import ModelConfig
-from nerfstudio.models.lerf import LERFModelConfig, LERFModel
+from typing import List, Literal, Tuple, Type
 
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+from nerfstudio.configs import base_config as cfg
+from nerfstudio.data.datamanagers.base_datamanager import (
+    DataManagerConfig,
+    VanillaDataManagerConfig,
+)
+from nerfstudio.data.datamanagers.lerf_datamanager import (
+    LERFDataManager,
+    LERFDataManagerConfig,
+)
+from nerfstudio.models.base_model import ModelConfig
+from nerfstudio.models.lerf import LERFModel, LERFModelConfig
+from nerfstudio.pipelines.base_pipeline import (
+    Pipeline,
+    VanillaPipeline,
+    VanillaPipelineConfig,
+)
+
 try:
     import open_clip
 except ImportError:
@@ -23,9 +29,11 @@ except ImportError:
 
 import torch
 
+
 @dataclass
 class OpenCLIPNetworkConfig(cfg.InstantiateConfig):
     """Configuration for network instantiation"""
+
     _target: Type = field(default_factory=lambda: OpenCLIPNetwork)
     """target class to instantiate"""
     clip_model_type: str = "ViT-B-16"
@@ -33,20 +41,21 @@ class OpenCLIPNetworkConfig(cfg.InstantiateConfig):
     clip_n_dims: int = 512
     negatives: Tuple[str] = ("object", "things", "stuff", "texture")
 
+
 # ambivalent about this being here, it's still technically rendering but I want to keep it with the OpenCLIPNetworkConfig
 class OpenCLIPNetwork:
     def __init__(self, config: OpenCLIPNetworkConfig):
         self.config = config
         model, _, _ = open_clip.create_model_and_transforms(
-            self.config.clip_model_type, # e.g., ViT-B-16
-            pretrained=self.config.clip_model_pretrained, # e.g., laion2b_s34b_b88k
+            self.config.clip_model_type,  # e.g., ViT-B-16
+            pretrained=self.config.clip_model_pretrained,  # e.g., laion2b_s34b_b88k
             precision="fp16",
         )
         self.tokenizer = open_clip.get_tokenizer(self.config.clip_model_type)
         self.model = model.to("cuda")
         self.clip_n_dims = self.config.clip_n_dims
 
-        self.positives = ["hand sanitizer"]
+        self.positives = ["'wes' text"]
         self.negatives = self.config.negatives
         with torch.no_grad():
             tok_phrases = torch.cat([self.tokenizer(phrase) for phrase in self.positives]).to("cuda")
@@ -56,8 +65,12 @@ class OpenCLIPNetwork:
         self.pos_embeds /= self.pos_embeds.norm(dim=-1, keepdim=True)
         self.neg_embeds /= self.neg_embeds.norm(dim=-1, keepdim=True)
 
-        assert self.pos_embeds.shape[1] == self.neg_embeds.shape[1], "Positive and negative embeddings must have the same dimensionality"
-        assert self.pos_embeds.shape[1] == self.clip_n_dims, "Embedding dimensionality must match the model dimensionality"
+        assert (
+            self.pos_embeds.shape[1] == self.neg_embeds.shape[1]
+        ), "Positive and negative embeddings must have the same dimensionality"
+        assert (
+            self.pos_embeds.shape[1] == self.clip_n_dims
+        ), "Embedding dimensionality must match the model dimensionality"
 
     def set_positives(self, text_list):
         self.positives = text_list
@@ -77,7 +90,9 @@ class OpenCLIPNetwork:
         sims = torch.stack((repeated_pos, negative_vals), dim=-1)  # rays x N-phrase x 2
         softmax = torch.softmax(10 * sims, dim=-1)  # rays x n-phrase x 2
         best_id = softmax[..., 0].argmin(dim=1)  # rays x 2
-        return torch.gather(softmax, 1, best_id[..., None, None].expand(best_id.shape[0], len(self.negatives), 2))[:, 0, :]
+        return torch.gather(softmax, 1, best_id[..., None, None].expand(best_id.shape[0], len(self.negatives), 2))[
+            :, 0, :
+        ]
 
 
 @dataclass
@@ -102,7 +117,7 @@ class LERFPipeline(VanillaPipeline):
         test_mode: Literal["test", "val", "inference"] = "val",
         world_size: int = 1,
         local_rank: int = 0,
-        ):
+    ):
         super(VanillaPipeline, self).__init__()
         self.config = config
         self.test_mode = test_mode
@@ -128,4 +143,4 @@ class LERFPipeline(VanillaPipeline):
         self.world_size = world_size
         if world_size > 1:
             self._model = typing.cast(LERFModel, DDP(self._model, device_ids=[local_rank], find_unused_parameters=True))
-            dist.barrier(device_ids=[local_rank])
+            dist.barrier(device_ids=[local_rank])