From 71476e1923409e2d4183e90af082e6be9b57860a Mon Sep 17 00:00:00 2001 From: Justin Date: Fri, 24 Mar 2023 19:56:23 +0000 Subject: [PATCH] downsample lerf samples --- nerfstudio/configs/method_configs.py | 8 ++-- nerfstudio/models/lerf.py | 36 ++++++++++++++---- nerfstudio/pipelines/lerf_pipeline.py | 55 +++++++++++++++++---------- 3 files changed, 67 insertions(+), 32 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 8cd4db38a..e2e3707f2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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), diff --git a/nerfstudio/models/lerf.py b/nerfstudio/models/lerf.py index 30d197a37..784600b0b 100644 --- a/nerfstudio/models/lerf.py +++ b/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 diff --git a/nerfstudio/pipelines/lerf_pipeline.py b/nerfstudio/pipelines/lerf_pipeline.py index ac6ea9b53..46b087667 100644 --- a/nerfstudio/pipelines/lerf_pipeline.py +++ b/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]) \ No newline at end of file + dist.barrier(device_ids=[local_rank])