From 151ed5c86796048c2379669c9cb4f9eb9f4358d1 Mon Sep 17 00:00:00 2001 From: Justin Date: Thu, 23 Mar 2023 19:43:49 +0000 Subject: [PATCH] small changes, move to diff optimizers again --- .../data/datamanagers/lerf_datamanager.py | 5 -- nerfstudio/fields/lerf_field.py | 24 ++++---- nerfstudio/models/lerf.py | 60 +++++++++---------- 3 files changed, 41 insertions(+), 48 deletions(-) diff --git a/nerfstudio/data/datamanagers/lerf_datamanager.py b/nerfstudio/data/datamanagers/lerf_datamanager.py index 3ef193668..c9419e82c 100644 --- a/nerfstudio/data/datamanagers/lerf_datamanager.py +++ b/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""" diff --git a/nerfstudio/fields/lerf_field.py b/nerfstudio/fields/lerf_field.py index 7ecee87ff..358178a42 100644 --- a/nerfstudio/fields/lerf_field.py +++ b/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 \ No newline at end of file + return output diff --git a/nerfstudio/models/lerf.py b/nerfstudio/models/lerf.py index 125d34eb4..ef83f352f 100644 --- a/nerfstudio/models/lerf.py +++ b/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 -