small changes, move to diff optimizers again

nerfstudio/data/datamanagers/lerf_datamanager.py

@@ -198,14 +198,9 @@ class LERFDataManager(VanillaDataManager):  # pylint: disable=abstract-method
             },
             cache_dir=Path(f"outputs/{self.config.dataparser.data.name}"),
         )
-        # self.clip_interpolator.try_load(f"outputs/{self.config.dataparser.data.name}")
-        # self.clip_interpolator.create()
-        # self.clip_interpolator.save(f"outputs/{self.config.dataparser.data.name}")
         self.dino_dataloader = DinoDataloader(
             image_list=images, device=self.device, cfg={}, cache_dir=Path(f"outputs/{self.config.dataparser.data.name}")
         )
-        # self.dino_dataloader.create()
-        # self.dino_dataloader.save(f"outputs/{self.config.dataparser.data.name}")
 
     def setup_eval(self):
         """Sets up the data loader for evaluation"""

nerfstudio/fields/lerf_field.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple, List
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -9,15 +9,12 @@ from torchtyping import TensorType
 from nerfstudio.cameras.rays import RaySamples
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.field_components.activations import trunc_exp
-from nerfstudio.field_components.field_heads import (
-    FieldHeadNames,
-)
-from nerfstudio.fields.base_field import Field
-
+from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import (
     SceneContraction,
     SpatialDistortion,
 )
+from nerfstudio.fields.base_field import Field
 
 try:
     import tinycudann as tcnn
@@ -30,10 +27,9 @@ class LERFField(Field):
         super().__init__()
 
         self.spatial_distortion = spatial_distortion
-        self.clip_encs = torch.nn.ModuleList([
-            LERFField._get_encoding(16, 128, 16, indim=3), 
-            LERFField._get_encoding(128, 512, 16, indim=3)
-            ])
+        self.clip_encs = torch.nn.ModuleList(
+            [LERFField._get_encoding(16, 128, 16, indim=3), LERFField._get_encoding(128, 512, 16, indim=3)]
+        )
         tot_out_dims = sum([e.n_output_dims for e in self.clip_encs])
 
         self.clip_net = tcnn.Network(
@@ -99,8 +95,10 @@ class LERFField(Field):
 
     def get_output_from_hashgrid(self, ray_samples: RaySamples, hashgrid_field, scale):
         # designated scales, run outputs for each scale
-        hashgrid_field = hashgrid_field.view(-1, self.clip_net.n_input_dims-1)
-        clip_pass = self.clip_net(torch.cat([hashgrid_field, scale.view(-1, 1)], dim=-1)).view(*ray_samples.frustums.shape, -1)
+        hashgrid_field = hashgrid_field.view(-1, self.clip_net.n_input_dims - 1)
+        clip_pass = self.clip_net(torch.cat([hashgrid_field, scale.view(-1, 1)], dim=-1)).view(
+            *ray_samples.frustums.shape, -1
+        )
         output = clip_pass / clip_pass.norm(dim=-1, keepdim=True)
 
-        return output
+        return output

nerfstudio/models/lerf.py

@@ -1,33 +1,35 @@
-import torch
+from dataclasses import dataclass, field
 from typing import Dict, List, Tuple, Type
-from torch.nn import Parameter
 
 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.renderers import (
+    CLIPRenderer,
+    DepthRenderer,
+    MeanRenderer,
+)
 from nerfstudio.models.nerfacto import NerfactoModel, NerfactoModelConfig
-from nerfstudio.model_components.renderers import CLIPRenderer, MeanRenderer, DepthRenderer
-from nerfstudio.cameras.rays import RayBundle
 
-from typing import Type
-from nerfstudio.data.scene_box import SceneBox
-from dataclasses import dataclass, field
 
 @dataclass
 class LERFModelConfig(NerfactoModelConfig):
     _target: Type = field(default_factory=lambda: LERFModel)
-    clip_loss_weight: float = 1
+    clip_loss_weight: float = 0.1
     n_scales: int = 30
+    max_scale: float = 1.5
+    """maximum scale used to compute relevancy with"""
     specify_scale: bool = False
 
+
 class LERFModel(NerfactoModel):
     config: LERFModelConfig
 
@@ -37,38 +39,39 @@ class LERFModel(NerfactoModel):
         self.renderer_clip = CLIPRenderer()
         self.renderer_mean = MeanRenderer()
 
-        self.network = self.kwargs['network']
+        self.network = self.kwargs["network"]
         self.lerf_field = LERFField(clip_n_dims=self.network.clip_n_dims)
 
-    def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, preset_scales = []):
+    def get_max_across(self, ray_samples, weights, hashgrid_field, scales_shape, preset_scales=None):
         # TODO smoothen this out
         if self.config.specify_scale:
-            assert len(preset_scales) != 0
+            assert preset_scales is not None
             assert len(preset_scales) == len(self.network.positives)
             scales_list = torch.tensor(preset_scales)
         else:
-            scales_list = torch.linspace(0.0, 1.0, self.config.n_scales)
+            scales_list = torch.linspace(0.0, self.config.max_scale, self.config.n_scales)
 
         # probably not a good idea bc it's prob going to be a lot of memory
         n_phrases = len(self.network.positives)
         n_phrases_maxs = [None for _ in range(n_phrases)]
         n_phrases_sims = [None for _ in range(n_phrases)]
-        scale_tensor = torch.ones(scales_shape, device=weights.device)
         for s, scale in enumerate(scales_list):
             scale = scale.item()
             with torch.no_grad():
-                clip_output = self.lerf_field.get_output_from_hashgrid(ray_samples, hashgrid_field, scale*scale_tensor)
+                clip_output = self.lerf_field.get_output_from_hashgrid(
+                    ray_samples,
+                    hashgrid_field,
+                    torch.full(scales_shape, s, device=weights.device, dtype=hashgrid_field.dtype),
+                )
             clip_output = self.renderer_clip(embeds=clip_output, weights=weights.detach())
 
             for i in range(n_phrases):
                 probs = self.network.get_relevancy(clip_output, i)
                 pos_prob = probs[..., 0:1]
-                print(scale, pos_prob.min(), pos_prob.max())
                 if n_phrases_maxs[i] is None or pos_prob.max() > n_phrases_sims[i].max():
                     n_phrases_maxs[i] = s
                     n_phrases_sims[i] = pos_prob
-        
-        print(n_phrases_maxs, scales_list, n_phrases_sims[0].min(), n_phrases_sims[0].max())
+
         return torch.stack(n_phrases_sims), torch.Tensor(n_phrases_maxs)
 
     def get_outputs(self, ray_bundle: RayBundle):
@@ -79,9 +82,7 @@ class LERFModel(NerfactoModel):
             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)
-            clip_scales = (
-                clip_scales * ray_bundle.metadata['width'] * (1/ray_bundle.metadata['fx']) * dist
-            )
+            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)
 
@@ -94,17 +95,17 @@ class LERFModel(NerfactoModel):
             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())
+        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())
 
         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, [0.111]
+                ray_samples, weights, lerf_field_outputs[FieldHeadNames.HASHGRID], clip_scales.shape
             )
             multiphrase = max_across[0]
-            print(max_across[0].min(), max_across[0].max())
             # normalization for sanity check TODO(cmk) remove
-            multiphrase[multiphrase < 0.5] = 0.5; multiphrase -= 0.5
+            multiphrase[multiphrase < 0.5] = 0.5
+            multiphrase -= 0.5
             multiphrase = multiphrase / multiphrase.max()
             outputs[f"multiphrase"] = multiphrase
 
@@ -133,9 +134,8 @@ class LERFModel(NerfactoModel):
         )
         loss_dict["dino_loss"] = torch.nn.functional.mse_loss(outputs["dino"], batch["dino"])
         return loss_dict
-    
+
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = super().get_param_groups()
-        param_groups["fields"] = list(self.lerf_field.parameters())+list(self.field.parameters())
+        param_groups["lerf"] = list(self.lerf_field.parameters())
         return param_groups
-