From 8fd7ea13fe69d842a4882cedecb3885ad19022c1 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 23 Dec 2022 19:31:02 +0000 Subject: [PATCH] working on abllations --- .../data/datamanagers/base_datamanager.py | 12 ++- nerfstudio/engine/trainer.py | 4 - nerfstudio/models/nerfacto.py | 21 ++++- nerfstudio/pipelines/base_pipeline.py | 12 +-- pyproject.toml | 1 + scripts/benchmarking/benchmark_nerfacto.py | 82 ++++++++++++++++--- scripts/benchmarking/create_tables.py | 36 ++++++++ scripts/downloads/download_data.py | 19 +++++ 8 files changed, 156 insertions(+), 31 deletions(-) create mode 100644 scripts/benchmarking/create_tables.py diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 9f4a51824..555d618e4 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -18,6 +18,7 @@ Datamanager. from __future__ import annotations +import copy import random from abc import abstractmethod from dataclasses import dataclass, field @@ -382,7 +383,10 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch) - self.eval_camera_optimizer = self.config.camera_optimizer.setup( + # eval camera optimizer should always have pose opt on + eval_camera_optimizer = copy.deepcopy(self.config.camera_optimizer) + eval_camera_optimizer.mode = "SO3xR3" + self.eval_camera_optimizer = eval_camera_optimizer.setup( num_cameras=self.eval_dataset.cameras.size, device=self.device ) self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device), self.eval_camera_optimizer) @@ -450,11 +454,11 @@ class VanillaDataManager(DataManager): # pylint: disable=abstract-method camera_opt_params_eval = list(self.eval_camera_optimizer.parameters()) if self.config.camera_optimizer.mode != "off": assert len(camera_opt_params_train) > 0 - assert len(camera_opt_params_eval) > 0 param_groups[self.config.camera_optimizer.param_group + "_train"] = camera_opt_params_train - param_groups[self.config.camera_optimizer.param_group + "_eval"] = camera_opt_params_eval else: assert len(camera_opt_params_train) == 0 - assert len(camera_opt_params_eval) == 0 + + assert len(camera_opt_params_eval) > 0 + param_groups[self.config.camera_optimizer.param_group + "_eval"] = camera_opt_params_eval return param_groups diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 0221316fe..b5089df1e 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -148,10 +148,6 @@ class Trainer: "optimizer": camera_optimizer_config.optimizer, "scheduler": camera_optimizer_config.scheduler, } - # optimizer_config[camera_optimizer_config.param_group + "_eval"] = { - # "optimizer": camera_optimizer_config.optimizer, - # "scheduler": camera_optimizer_config.scheduler, - # } return Optimizers(optimizer_config, param_groups) def train(self) -> None: diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 2e46b657e..b531cff54 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -67,7 +67,7 @@ class NerfactoModelConfig(ModelConfig): """How far along the ray to start sampling.""" far_plane: float = 1000.0 """How far along the ray to stop sampling.""" - background_color: Literal["background", "last_sample"] = "last_sample" + background_color: Literal["random", "last_sample"] = "last_sample" """Whether to randomize the background color.""" num_levels: int = 16 """Number of levels of the hashmap for the base mlp.""" @@ -116,6 +116,10 @@ class NerfactoModelConfig(ModelConfig): """Whether use single jitter or not for the proposal networks.""" predict_normals: bool = False """Whether to predict normals or not.""" + use_scene_contraction: bool = True + """Whether to use scene contraction or not.""" + scene_contraction_norm: Literal["inf", "l2"] = "inf" + """Which norm to use for the scene contraction.""" class NerfactoModel(Model): @@ -131,7 +135,15 @@ class NerfactoModel(Model): """Set the fields and modules.""" super().populate_modules() - scene_contraction = SceneContraction(order=float("inf")) + # TODO(ethan): flag to enable scene contraction or not + + if self.config.scene_contraction_norm == "inf": + order = float("inf") + elif self.config.scene_contraction_norm == "l2": + order = None + else: + raise ValueError("Invalid scene contraction norm") + scene_contraction = SceneContraction(order=order) # Fields self.field = TCNNNerfactoField( @@ -152,8 +164,9 @@ class NerfactoModel(Model): # Build the proposal network(s) self.proposal_networks = torch.nn.ModuleList() if self.config.use_same_proposal_network: - assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." - prop_net_args = self.config.proposal_net_args_list[0] + # assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed." + # using the last defined proposal network + prop_net_args = self.config.proposal_net_args_list[-1] network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args) self.proposal_networks.append(network) self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)]) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index ba46e9c4b..778458910 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -203,7 +203,7 @@ class VanillaPipelineConfig(cfg.InstantiateConfig): """Whether to optimize the appearance during evaluation.""" eval_num_appearance_iters: int = 10 """Number of iterations to optimize the appearance.""" - eval_image_scale_factor: float = 0.25 + eval_image_scale_factor: float = 0.125 """Scale factor to use for evaluation images so that they fit in memory.""" @@ -367,11 +367,11 @@ class VanillaPipeline(Pipeline): param_groups = self.datamanager.get_param_groups() optimizer = torch.optim.Adam(param_groups[camera_opt_param_group], lr=1e-3, eps=1e-15) - # check that camera optimizer is on - if self.config.datamanager.camera_optimizer.mode == "off": - raise ValueError( - "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses." - ) + # # check that camera optimizer is on + # if self.config.datamanager.camera_optimizer.mode == "off": + # raise ValueError( + # "You must set datamanager.camera_optimizer.mode to optimize to optimize the camera poses." + # ) # CONSOLE.print("Optimizing camera poses...") # rescale the image to 1/4 resolution diff --git a/pyproject.toml b/pyproject.toml index fe908f8c6..b35bf330f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ dependencies = [ "av==9.2.0", "tyro>=0.3.31", "gdown==4.5.1", + "GPUtil>=1.4.0", "ninja==1.10.2.3", "functorch==0.2.1", "h5py>=2.9.0", diff --git a/scripts/benchmarking/benchmark_nerfacto.py b/scripts/benchmarking/benchmark_nerfacto.py index 082590cad..e450217f0 100644 --- a/scripts/benchmarking/benchmark_nerfacto.py +++ b/scripts/benchmarking/benchmark_nerfacto.py @@ -4,6 +4,11 @@ Process the Mipnerf360 dataset. from nerfstudio.utils.scripts import run_command +import GPUtil +import time +import threading + + capture_names = [ "Egypt", "person", @@ -17,36 +22,87 @@ capture_names = [ "Giannini-Hall", ] table_rows = [ - ("nerfacto", ""), - ("w/o-pose", "--pipeline.datamanager.camera-optimizer.mode off"), - ("w/o-app", "--pipeline.model.use-appearance-embedding False"), + ("nerfacto", "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True"), + ( + "w/o-pose", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.datamanager.camera-optimizer.mode off", + ), + ( + "w/o-app", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.model.use-appearance-embedding False", + ), ( "w/o-pose-app", - "--pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance False --pipeline.datamanager.camera-optimizer.mode off --pipeline.model.use-appearance-embedding False", ), ( "1-prop-network", - "--pipeline.model.num-proposal-samples-per-ray \"256\" --pipeline.model.num_proposal_iterations 1", + '--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.num-proposal-samples-per-ray "256" --pipeline.model.num_proposal_iterations 1', + ), + ( + "l2-contraction", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.scene-contraction-norm l2", + ), + ( + "shared-prop-network", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.use-same-proposal-network True", + ), + ( + "random-background-color", + "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True --pipeline.model.background-color random", ), ] # 30K iterations +# make a list of all the jobs that need to be fun +jobs = [] for capture_name in capture_names: for table_row_name, table_row_command in table_rows: - - # TODO: find an available GPU - gpu = 0 - command = " ".join( ( - # f"CUDA_VISIBLE_DEVICES={gpu}", f"ns-train nerfacto", "--vis wandb", f"--data data/nerfstudio/{capture_name}", - "--pipeline.eval_optimize_cameras True --pipeline.eval_optimize_appearance True", - f"--wandb-name {capture_name}-{table_row_name} --experiment-name {capture_name}-{table_row_name}", + "--output-dir outputs/nerfacto-ablations", + "--trainer.steps-per-eval-batch 0 --trainer.steps-per-eval-image 0", + "--trainer.steps-per-eval-all-images 5000 --trainer.max-num-iterations 30001", + f"--wandb-name {capture_name}_{table_row_name}", + f"--experiment-name {capture_name}_{table_row_name}", table_row_command, ) ) - print(command) + jobs.append(command) + +while jobs: + + # check which GPUs have capacity to run these jobs + """Returns the available GPUs.""" + gpu_devices_available = GPUtil.getAvailable(order="first", limit=10, maxMemory=0.1) + + print("Available GPUs: ", gpu_devices_available) + + # thread list + threads = [] + while gpu_devices_available and jobs: + gpu = gpu_devices_available.pop(0) + command = f"CUDA_VISIBLE_DEVICES={gpu} " + jobs.pop(0) + + def task(): + print("Starting command: ", command) + out = run_command(command, verbose=False) + # time.sleep(5) + print("Finished command: ", command) + + threads.append(threading.Thread(target=task)) + threads[-1].start() + + # NOTE(ethan): need a delay otherwise the wandb/tensorboard naming is messed up + # not sure why? + time.sleep(5) + + # wait for all threads to finish + for t in threads: + t.join() + + print("Finished all threads") diff --git a/scripts/benchmarking/create_tables.py b/scripts/benchmarking/create_tables.py new file mode 100644 index 000000000..25b361b24 --- /dev/null +++ b/scripts/benchmarking/create_tables.py @@ -0,0 +1,36 @@ +""" +Code to create the tables. +""" + +import pandas as pd + +# from nerfstudio.scripts.benchmarking.benchmark_nerfacto import table_rows + +# the csv filenames +csv_filename = "wandb-psnr.csv" + +# code to read the csv file + +df = pd.read_csv(csv_filename) + +column_titles = list(df.columns.values) +last_row = list(df.values[-1]) + +print(len(column_titles)) +print(len(last_row)) + +from collections import defaultdict +results = defaultdict(dict) +for column_title, last_row_value in zip(column_titles, last_row): + if column_title == "Step" or "__MAX" in column_title or "__MIN" in column_title: + continue + name = column_title.split(" - ")[0] + capture_name, method = name.split("_") + metric = float(last_row_value) + # format to have 2 decimal places + metric = f"{metric:.2f}" + results[capture_name][method] = metric + +import pprint +pprint.pprint(results) +print(len(results)) diff --git a/scripts/downloads/download_data.py b/scripts/downloads/download_data.py index 12a0a0227..9636d0b58 100755 --- a/scripts/downloads/download_data.py +++ b/scripts/downloads/download_data.py @@ -97,6 +97,15 @@ nerfstudio_dataset = [ "sculpture", "Giannini-Hall", ] +mipnerf360_dataset = [ + "mipnerf-360-bicycle", + "mipnerf-360-garden", + "mipnerf-360-stump", + "mipnerf-360-room", + "mipnerf-360-counter", + "mipnerf-360-kitchen", + "mipnerf-360-bonsai", +] nerfstudio_file_ids = { "bww_entrance": grab_file_id("https://drive.google.com/file/d/1ylkRHtfB3n3IRLf2wplpfxzPTq7nES9I/view?usp=sharing"), "campanile": grab_file_id("https://drive.google.com/file/d/13aOfGJRRH05pOOk9ikYGTwqFc2L1xskU/view?usp=sharing"), @@ -106,6 +115,7 @@ nerfstudio_file_ids = { "redwoods2": grab_file_id("https://drive.google.com/file/d/1rg-4NoXT8p6vkmbWxMOY6PSG4j3rfcJ8/view?usp=sharing"), "storefront": grab_file_id("https://drive.google.com/file/d/16b792AguPZWDA_YC4igKCwXJqW0Tb21o/view?usp=sharing"), "vegetation": grab_file_id("https://drive.google.com/file/d/1wBhLQ2odycrtU39y2akVurXEAt9SsVI3/view?usp=sharing"), + # Nerfstudio Dataset "Egypt": "https://data.nerf.studio/nerfstudio/Egypt.zip", "person": "https://data.nerf.studio/nerfstudio/person.zip", "kitchen": "https://data.nerf.studio/nerfstudio/kitchen.zip", @@ -116,8 +126,17 @@ nerfstudio_file_ids = { "stump": "https://data.nerf.studio/nerfstudio/stump.zip", "sculpture": "https://data.nerf.studio/nerfstudio/sculpture.zip", "Giannini-Hall": "https://data.nerf.studio/nerfstudio/Giannini-Hall.zip", + # MipNeRF-360 Dataset. TODO: add these links + "mipnerf360-bicycle": "", + "mipnerf360-garden": "", + "mipnerf360-stump": "", + "mipnerf360-room": "", + "mipnerf360-counter": "", + "mipnerf360-kitchen": "", + "mipnerf360-bonsai": "", "all": None, "nerfstudio-dataset": nerfstudio_dataset, + "mipnerf360-dataset": mipnerf360_dataset, } NerfstudioCaptureName = tyro.extras.literal_type_from_choices(nerfstudio_file_ids.keys())