From 8973039e6519082cdf64029e40bdd5e3c3ead038 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 15 Sep 2022 00:16:41 -0700 Subject: [PATCH] misc cleanup (#432) * misc cuda improvements. added cone angle as param * wandb as default. better organization for experiments * turn off tensorboard * disable eval when using viewer * remove comment * launch.json --- .vscode/launch.json | 19 +++++ nerfactory/configs/base.py | 44 ++++++----- nerfactory/configs/base_configs.py | 4 +- nerfactory/cuda/csrc/intersection.cu | 10 +-- nerfactory/cuda/csrc/raymarching.cu | 35 ++++++--- nerfactory/engine/trainer.py | 11 ++- nerfactory/fields/instant_ngp_field.py | 14 +++- nerfactory/models/base.py | 4 +- nerfactory/models/compound.py | 3 +- nerfactory/models/instant_ngp.py | 93 ++++++++++++++++-------- nerfactory/models/mipnerf.py | 2 +- nerfactory/models/mipnerf_360.py | 2 +- nerfactory/models/modules/ray_sampler.py | 9 +-- nerfactory/models/nerfw.py | 6 ++ nerfactory/models/tensorf.py | 4 +- nerfactory/models/vanilla_nerf.py | 2 +- nerfactory/pipelines/base.py | 6 +- nerfactory/utils/writer.py | 64 ++++++++++++---- pyproject.toml | 2 +- scripts/run_eval.py | 3 +- tests/test_run_train.py | 4 + 21 files changed, 230 insertions(+), 111 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 5cbcd1a53..c7d3f4781 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -49,6 +49,25 @@ "instant_ngp" ] }, + { + "name": "Python: run_train ngp nerfactory (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "console": "integratedTerminal", + "args": [ + "instant_ngp", + "--experiment_name=nerfactory_experiment", + "--logging.writer.2.max-log-size=0", + // viewer + "--viewer.enable", + "--viewer.zmq-url=tcp://127.0.0.1:6001", + "--viewer.websocket-port=8014", + "--pipeline.model.randomize-background", + "--pipeline.model.cone-angle=0.00390625", + "nerfactory-data" + ] + }, { "name": "Python: run_train ngp-friends (gpu)", "type": "python", diff --git a/nerfactory/configs/base.py b/nerfactory/configs/base.py index 034942361..6bedec042 100644 --- a/nerfactory/configs/base.py +++ b/nerfactory/configs/base.py @@ -23,6 +23,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Type import dcargs import torch +from rich import console from nerfactory.configs.utils import to_immutable_dict @@ -39,13 +40,13 @@ from nerfactory.datamanagers.dataparsers.record3d_parser import Record3D # model instances from nerfactory.models.base import Model from nerfactory.models.compound import CompoundModel -from nerfactory.models.instant_ngp import NGPModel from nerfactory.models.nerfw import NerfWModel from nerfactory.models.tensorf import TensoRFModel from nerfactory.optimizers.schedulers import ExponentialDecaySchedule from nerfactory.pipelines.base import Pipeline from nerfactory.utils import writer +CONSOLE = console.Console() # Pretty printing class class PrintableConfig: # pylint: disable=too-few-public-methods @@ -116,7 +117,7 @@ class WandbWriterConfig(InstantiateConfig): """target class to instantiate""" enable: bool = False """if True enables wandb logging, else disables""" - relative_log_dir: Path = Path("./") + relative_log_dir: Path = Path("./wandb") """relative path to save all wandb events""" log_dir: dcargs.conf.Fixed[Path] = dcargs.MISSING """auto-populated absolute path to saved wandb events""" @@ -164,8 +165,8 @@ class LoggingConfig(PrintableConfig): """maximum history size to keep for computing running averages of stats. e.g. if 20, averages will be computed over past 20 occurances.""" writer: Tuple[Any, ...] = ( - TensorboardWriterConfig(enable=True), - WandbWriterConfig(enable=False), + TensorboardWriterConfig(enable=False), + WandbWriterConfig(enable=True), LocalWriterConfig(enable=True), ) """list of all supported writers. Can turn on/off writers by specifying enable.""" @@ -215,7 +216,7 @@ class NerfactoryDataParserConfig(DataParserConfig): _target: Type = Nerfactory """target class to instantiate""" - data_directory: Path = Path("data/ours/posterv2") + data_directory: Path = Path("data/ours/posters_v3") """directory specifying location of data""" scale_factor: float = 1.0 """How much to scale the camera origins by.""" @@ -369,7 +370,7 @@ class ModelConfig(InstantiateConfig): """target class to instantiate""" enable_collider: bool = True """Whether to create a scene collider to filter rays.""" - collider_params: Dict[str, float] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) + collider_params: Optional[Dict[str, float]] = to_immutable_dict({"near_plane": 2.0, "far_plane": 6.0}) """parameters to instantiate scene collider with""" loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss_coarse": 1.0, "rgb_loss_fine": 1.0}) """Loss specific weights.""" @@ -377,8 +378,6 @@ class ModelConfig(InstantiateConfig): """Number of samples in coarse field evaluation. Defaults to 64""" num_importance_samples: int = 128 """Number of samples in fine field evaluation. Defaults to 128""" - field_implementation: Literal["torch", "tcnn"] = "torch" - """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class""" enable_density_field: bool = False """Whether to create a density field to filter samples.""" density_field_params: Dict[str, Any] = to_immutable_dict( @@ -394,21 +393,8 @@ class ModelConfig(InstantiateConfig): @dataclass -class InstantNGPModelConfig(ModelConfig): - """Instant NGP Model Config""" - - _target: Type = NGPModel - """target class to instantiate""" - enable_collider: bool = False - """Whether to create a scene collider to filter rays.""" - loss_coefficients: Dict[str, float] = to_immutable_dict({"rgb_loss": 1.0}) - """Loss specific weights.""" - field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... - """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" - enable_density_field: bool = True - """Whether to create a density field to filter samples.""" - num_samples: int = 1024 # instead of course/fine samples - """Number of samples in field evaluation. Defaults to 1024,""" +class VanillaModelConfig(InstantiateConfig): + """Vanilla Model Config""" @dataclass @@ -560,6 +546,18 @@ class Config(PrintableConfig): """Make paths more specific using the current timestamp.""" self.set_timestamp() + # if the viewer is enabled, disable tensorboard and wandb logging + # TODO(ethan): this is gross for now and will require a config refactor + if self.viewer.enable: + string = "Disabling tensorboard and wandb logging since viewer is enabled." + CONSOLE.print(f"[bold red]{string}") + self.logging.writer[0].enable = False + self.logging.writer[1].enable = False + # also disable eval steps + string = "Disabling eval iterations since viewer is enabled." + CONSOLE.print(f"[bold red]{string}") + self.trainer.steps_per_test = self.trainer.max_num_iterations + def set_timestamp(self, timestamp: Optional[str] = None) -> None: """Make paths in our config more specific using a timestamp. diff --git a/nerfactory/configs/base_configs.py b/nerfactory/configs/base_configs.py index bccda6bc6..4734efbc8 100644 --- a/nerfactory/configs/base_configs.py +++ b/nerfactory/configs/base_configs.py @@ -27,7 +27,6 @@ from nerfactory.configs.base import ( CompoundModelConfig, Config, FriendsDataManagerConfig, - InstantNGPModelConfig, MipNerf360DataParserConfig, ModelConfig, NerfWModelConfig, @@ -38,6 +37,7 @@ from nerfactory.configs.base import ( TrainerConfig, VanillaDataManagerConfig, ) +from nerfactory.models.instant_ngp import InstantNGPModelConfig from nerfactory.models.mipnerf import MipNerfModel from nerfactory.models.mipnerf_360 import MipNerf360Model from nerfactory.models.semantic_nerf import SemanticNerfModel @@ -46,7 +46,7 @@ from nerfactory.models.vanilla_nerf import NeRFModel base_configs: Dict[str, Config] = {} base_configs["instant_ngp"] = Config( method_name="instant_ngp", - trainer=TrainerConfig(mixed_precision=True), + trainer=TrainerConfig(steps_per_test=500, steps_per_save=2000, mixed_precision=True), pipeline=PipelineConfig( datamanager=VanillaDataManagerConfig( train_dataparser=BlenderDataParserConfig(), diff --git a/nerfactory/cuda/csrc/intersection.cu b/nerfactory/cuda/csrc/intersection.cu index 9b272c65f..d37abd81b 100644 --- a/nerfactory/cuda/csrc/intersection.cu +++ b/nerfactory/cuda/csrc/intersection.cu @@ -19,8 +19,8 @@ inline __host__ __device__ void _ray_aabb_intersect( if (tymin > tymax) __swap(tymin, tymax); if (tmin > tymax || tymin > tmax){ - *near = std::numeric_limits::max(); - *far = std::numeric_limits::max(); + *near = 1e10; + *far = 1e10; return; } @@ -32,8 +32,8 @@ inline __host__ __device__ void _ray_aabb_intersect( if (tzmin > tzmax) __swap(tzmin, tzmax); if (tmin > tzmax || tzmin > tmax){ - *near = std::numeric_limits::max(); - *far = std::numeric_limits::max(); + *near = 1e10; + *far = 1e10; return; } @@ -103,7 +103,7 @@ std::vector ray_aabb_intersect( AT_DISPATCH_FLOATING_TYPES_AND_HALF( rays_o.scalar_type(), "ray_aabb_intersect", ([&] { - kernel_ray_aabb_intersect<<>>( + kernel_ray_aabb_intersect<<>>( N, rays_o.data_ptr(), rays_d.data_ptr(), diff --git a/nerfactory/cuda/csrc/raymarching.cu b/nerfactory/cuda/csrc/raymarching.cu index 9cd5bfabe..c073ece12 100644 --- a/nerfactory/cuda/csrc/raymarching.cu +++ b/nerfactory/cuda/csrc/raymarching.cu @@ -8,6 +8,10 @@ inline __device__ float min_step_size(uint32_t max_ray_samples) { // Maximum step size is the width of the coarsest gridsize cell. inline __device__ float max_step_size(uint32_t grid_cascades, uint32_t grid_size) { // Note(ruilongli): use mip=0 step size? + // Example: + // if there are 3 grid cascades and the grid size is 128 resolution + // then we will have 2^(num_cascades-1) = 4 + // and sqrt(3) * 4 / 128 = 0.03125 return __SQRT3() * (1 << (grid_cascades-1)) / grid_size; } @@ -108,22 +112,22 @@ template __global__ void kernel_raymarching( // rays info const uint32_t n_rays, - const scalar_t* rays_o, - const scalar_t* rays_d, - const scalar_t* t_min, - const scalar_t* t_max, + const scalar_t* rays_o, // shape (n_rays, 3) + const scalar_t* rays_d, // shape (n_rays, 3) + const scalar_t* t_min, // shape (n_rays,) + const scalar_t* t_max, // shape (n_rays,) // density grid - const float grid_center, - const float grid_scale, - const int grid_cascades, - const int grid_size, + const float grid_center, // default 0.0 + const float grid_scale, // default 3.0 + const int grid_cascades, // default 1 + const int grid_size, // default 128 const uint8_t* grid_bitfield, // sampling const int marching_steps, // used only for inference. -1 means disable it. const int max_total_samples, - const int max_ray_samples, // default 1024 - const float cone_angle, // default 0. for nerf-syn and 1/256 for large scene - const float step_scale, + const int max_ray_samples, // default 1024 + const float cone_angle, // default 0. for nerf-syn and 1/256 for large scene + const float step_scale, // length of the longest side of the scene bounding box int* steps_counter, int* rays_counter, int* packed_info_out, @@ -148,7 +152,13 @@ __global__ void kernel_raymarching( // TODO(ruilongli): perturb `near` as in ngp_pl? // TODO(ruilongli): pre-compute `dt_min`, `dt_max` as it is a constant? float dt_min = min_step_size(max_ray_samples) * step_scale; - float dt_max = max_step_size(grid_cascades, grid_size) * grid_scale; + float dt_max = max_step_size(grid_cascades, grid_size) * grid_scale; + + // if (threadIdx.x == 0) { + // // these values should be equivalent with the current code + // printf("step_scale value %f\n", step_scale); + // printf("grid_scale value %f\n", grid_scale); + // } uint32_t ray_idx, base, marching_samples; uint32_t j; @@ -322,6 +332,7 @@ __global__ void kernel_occupancy_query( * @param max_ray_samples Used to define the minimal step size: SQRT3() / max_ray_samples. * @param cone_angle 0. for nerf-synthetic and 1./256 for real scenes. * @param step_scale Scale up the step size by this much. Usually equals to scene scale. + * For example, this is often set to the length of the longest side of the scene bounding box. * @return std::vector * - packed_info: Stores how to index the ray samples from the returned values. * Shape of [n_rays, 3]. First value is the ray index. Second value is the sample diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py index 4da8c07c6..ef6bbf1e7 100644 --- a/nerfactory/engine/trainer.py +++ b/nerfactory/engine/trainer.py @@ -17,6 +17,7 @@ Code to train model. """ from __future__ import annotations +import dataclasses import functools import logging import os @@ -93,6 +94,8 @@ class Trainer: ) profiler.setup_profiler(config.logging) + writer.put_config(name="config", config_dict=dataclasses.asdict(config), step=0) + def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions. @@ -147,12 +150,14 @@ class Trainer: "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing." ) - if step != 0 and step % self.config.logging.steps_per_log == 0: - writer.put_dict(name="Loss/train-loss_metrics_dict", scalar_dict=loss_metric_dict, step=step) + if step % self.config.logging.steps_per_log == 0: + writer.put_dict(name="Train Metrics and Loss", scalar_dict=loss_metric_dict, step=step) if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0: self._save_checkpoint(self.config.trainer.model_dir, step) if step % self.config.trainer.steps_per_test == 0: - self.pipeline.get_eval_loss_dict(step=step) + metrics_dict = self.pipeline.get_eval_loss_dict(step=step) + writer.put_dict(name="Eval Metrics", scalar_dict=metrics_dict, step=step) + # write out the image dictionary returned too self._update_rays_per_sec(train_t, vis_t, step) self._write_out_storage(step) diff --git a/nerfactory/fields/instant_ngp_field.py b/nerfactory/fields/instant_ngp_field.py index 6f79c928c..779563249 100644 --- a/nerfactory/fields/instant_ngp_field.py +++ b/nerfactory/fields/instant_ngp_field.py @@ -48,7 +48,7 @@ def get_normalized_directions(directions: TensorType["bs":..., 3]): class TCNNInstantNGPField(Field): - """NeRF Field + """TCNN implementation of the Instant-NGP field. Args: aabb: parameters of scene aabb bounds @@ -120,6 +120,11 @@ class TCNNInstantNGPField(Field): def get_density(self, ray_samples: RaySamples): positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) + # assert all positions are in the range [0, 1] + # otherwise print min and max values + # assert torch.all(positions >= 0.0) and torch.all( + # positions <= 1.0 + # ), f"positions: {positions.min()} {positions.max()}" h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) @@ -131,9 +136,12 @@ class TCNNInstantNGPField(Field): def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None): # TODO: add valid_mask masking! - # tcnn requires directions in the range [0,1] directions = get_normalized_directions(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) + # assert all directions are in the range [0, 1] + # assert torch.all(directions >= 0.0) and torch.all( + # directions <= 1.0 + # ), f"directions: {directions.min()} {directions.max()}" d = self.direction_encoding(directions_flat) if density_embedding is None: positions = SceneBounds.get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) @@ -146,7 +154,7 @@ class TCNNInstantNGPField(Field): class TorchInstantNGPField(NeRFField): """ - PyTorch implementation of the instant-ngp field. + PyTorch implementation of the Instant-NGP field. """ def __init__( diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index 3e9b65a21..8ec1acca0 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -207,7 +207,7 @@ class Model(nn.Module): return outputs @abstractmethod - def log_test_image_outputs(self, image_idx, step, batch, outputs) -> float: + def log_test_image_outputs(self, image_idx, step, batch, outputs) -> Dict[str, float]: """Writes the test image outputs. TODO: This shouldn't return a loss @@ -218,7 +218,7 @@ class Model(nn.Module): outputs: Outputs of the model. Returns: - The psnr. + A dictionary of metrics. """ def load_model(self, loaded_state: Dict[str, Any]) -> None: diff --git a/nerfactory/models/compound.py b/nerfactory/models/compound.py index 43e79d37b..ae9fdcc24 100644 --- a/nerfactory/models/compound.py +++ b/nerfactory/models/compound.py @@ -106,7 +106,6 @@ class CompoundModel(Model): param_groups["fields"] = list(self.field.parameters()) return param_groups - @torch.cuda.amp.autocast() def get_outputs(self, ray_bundle: RayBundle): # TODO(ruilongli) # - train test difference @@ -192,4 +191,4 @@ class CompoundModel(Model): writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step) - return psnr.item() + return {"psnr": float(psnr.item())} diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index 986e10a9d..b9357ec06 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -18,7 +18,8 @@ Implementation of Instant NGP. from __future__ import annotations -from typing import Dict, List +from dataclasses import dataclass, field +from typing import Dict, List, Literal, Optional, Type import torch from torch.nn import Parameter @@ -34,7 +35,7 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.models.base import Model from nerfactory.models.modules.ray_sampler import NGPSpacedSampler from nerfactory.optimizers.loss import MSELoss -from nerfactory.utils import colors, misc, visualization, writer +from nerfactory.utils import colors, visualization, writer from nerfactory.utils.callbacks import ( TrainingCallback, TrainingCallbackAttributes, @@ -42,6 +43,30 @@ from nerfactory.utils.callbacks import ( ) +@dataclass +class InstantNGPModelConfig(cfg.ModelConfig): + """Instant NGP Model Config""" + + _target: Type = field( + default_factory=lambda: NGPModel + ) # We can't write `NGPModel` directly, because `NGPModel` doesn't exist yet + """target class to instantiate""" + enable_collider: bool = False + """Whether to create a scene collider to filter rays.""" + collider_params: Optional[Dict[str, float]] = None + """Instant NGP doesn't use a collider.""" + field_implementation: Literal["torch", "tcnn"] = "tcnn" # torch, tcnn, ... + """one of "torch" or "tcnn", or other fields in 'field_implementation_to_class'""" + enable_density_field: bool = True + """Whether to create a density field to filter samples.""" + num_samples: int = 1024 + """Number of samples in field evaluation. Defaults to 1024,""" + cone_angle: float = 0.0 + """Should be set to 0.0 for blender scenes but 1./256 for real scenes.""" + randomize_background: bool = False + """Whether to randomize the background color. Defaults to False.""" + + class NGPModel(Model): """Instant NGP model @@ -49,9 +74,9 @@ class NGPModel(Model): config: instant NGP configuration to instantiate model """ - config: cfg.InstantNGPModelConfig + config: InstantNGPModelConfig - def __init__(self, config: cfg.InstantNGPModelConfig, **kwargs) -> None: + def __init__(self, config: InstantNGPModelConfig, **kwargs) -> None: assert config.field_implementation in field_implementation_to_class self.field = None super().__init__(config=config, **kwargs) @@ -86,8 +111,6 @@ class NGPModel(Model): self.ssim = structural_similarity_index_measure self.lpips = LearnedPerceptualImagePatchSimilarity() - # no colliders default - def get_param_groups(self) -> Dict[str, List[Parameter]]: param_groups = {} if self.field is None: @@ -95,17 +118,18 @@ class NGPModel(Model): param_groups["fields"] = list(self.field.parameters()) return param_groups - @torch.cuda.amp.autocast() def get_outputs(self, ray_bundle: RayBundle): + assert self.field is not None + # TODO(ruilongli) # - train test difference # - visualize "depth_density_grid" num_rays = len(ray_bundle) device = ray_bundle.origins.device - if self.field is None: - raise ValueError("populate_fields() must be called before get_outputs") - ray_samples, packed_info, t_min, t_max = self.sampler(ray_bundle, self.field.aabb) + ray_samples, packed_info, t_min, t_max = self.sampler( + ray_bundle, self.field.aabb, cone_angle=self.config.cone_angle + ) field_outputs = self.field.forward(ray_samples) rgbs = field_outputs[FieldHeadNames.RGB] @@ -127,7 +151,13 @@ class NGPModel(Model): opacities, ) accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None]) - accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight) + + if self.config.randomize_background: + background_colors = torch.rand((num_rays, 3), device=device) # (num_rays, 3) + else: + background_colors = colors.WHITE.to(accumulated_color) # (3,) + + accumulated_color = accumulated_color + background_colors * (1.0 - accumulated_weight) outputs = { "rgb": accumulated_color, @@ -138,33 +168,32 @@ class NGPModel(Model): return outputs def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): - device = self.device - image = batch["image"].to(device) - if "alive_ray_mask" in outputs: - mask = outputs["alive_ray_mask"] - rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) - else: - rgb_loss = self.rgb_loss(image, outputs["rgb"]) + image = batch["image"].to(self.device) + mask = outputs["alive_ray_mask"] + rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) loss_dict = {"rgb_loss": rgb_loss} - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): - image = batch["image"].to(outputs["rgb"].device) + image = batch["image"].to(self.device) rgb = outputs["rgb"] acc = visualization.apply_colormap(outputs["accumulation"]) depth = visualization.apply_depth_colormap( outputs["depth"], accumulation=outputs["accumulation"], ) + alive_ray_mask = visualization.apply_colormap(outputs["alive_ray_mask"]) combined_rgb = torch.cat([image, rgb], dim=1) combined_acc = torch.cat([acc], dim=1) combined_depth = torch.cat([depth], dim=1) + combined_alive_ray_mask = torch.cat([alive_ray_mask], dim=1) - writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step) - writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) - writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) + # TODO: return an image dictionary instead of logging here + writer.put_image(name="Eval Images/img", image=combined_rgb, step=step) + writer.put_image(name="Eval Images/accumulation", image=combined_acc, step=step) + writer.put_image(name="Eval Images/depth", image=combined_depth, step=step) + writer.put_image(name="Eval Images/alive_ray_mask", image=combined_alive_ray_mask, step=step) # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] @@ -174,10 +203,14 @@ class NGPModel(Model): ssim = self.ssim(image, rgb) lpips = self.lpips(image, rgb) - writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(psnr), step=step) - writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step) # type: ignore - writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step) - - writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step) - - return psnr.item() + # all of these metrics will be logged as scalars + metrics_dict = { + "psnr": float(psnr.item()), + "image_idx": image_idx, + "ssim": float(ssim), # type: ignore + "lpips": float(lpips), + "acc_min": float(acc.min().item()), + "acc_max": float(acc.max().item()), + } + # TODO(ethan): return an image dictionary + return metrics_dict diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index 1e7772c76..125bf8a85 100644 --- a/nerfactory/models/mipnerf.py +++ b/nerfactory/models/mipnerf.py @@ -191,4 +191,4 @@ class MipNerfModel(Model): writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) - return fine_psnr.item() + return {"psnr": float(fine_psnr.item())} diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index f9807efac..4ad3c0b2a 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/nerfactory/models/mipnerf_360.py @@ -210,4 +210,4 @@ class MipNerf360Model(Model): writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) - return fine_psnr.item() + return {"psnr": float(fine_psnr.item())} diff --git a/nerfactory/models/modules/ray_sampler.py b/nerfactory/models/modules/ray_sampler.py index 5ca58c83a..bcccfb4ba 100644 --- a/nerfactory/models/modules/ray_sampler.py +++ b/nerfactory/models/modules/ray_sampler.py @@ -360,6 +360,7 @@ class NGPSpacedSampler(Sampler): ray_bundle: RayBundle, aabb: TensorType[2, 3], num_samples: Optional[int] = None, + cone_angle: Optional[float] = 0.0, marching_steps: Optional[int] = 128, t_min: Optional[TensorType["num_rays"]] = None, t_max: Optional[TensorType["num_rays"]] = None, @@ -373,6 +374,7 @@ class NGPSpacedSampler(Sampler): ray_bundle: Rays to generate samples for aabb: Bounding box of the scene. num_samples: Number of samples per ray + cone_angle: 0.0 for nerf-synthetic and 1.0/256 for real scenes Returns: First return is ray samples in a packed way where only the valid samples are kept. @@ -385,16 +387,13 @@ class NGPSpacedSampler(Sampler): num_samples = num_samples or self.num_samples - aabb = aabb.flatten() + aabb = aabb.flatten().contiguous() rays_o = ray_bundle.origins.contiguous() rays_d = ray_bundle.directions.contiguous() scene_scale = (aabb[3:6] - aabb[0:3]).max() if t_min is None or t_max is None: - # TODO(ruilongli): this clipping here is stupid. Try to avoid that. t_min, t_max = nerfactory_cuda.ray_aabb_intersect(rays_o, rays_d, aabb) - t_min = torch.clamp(t_min, max=1e10) # type: ignore - t_max = torch.clamp(t_max, max=1e10) # type: ignore marching_steps = -1 # disable marching mode for now # TODO(ruilongli): * 16 is for original impl for training. Need to run @@ -418,7 +417,7 @@ class NGPSpacedSampler(Sampler): marching_steps, max_samples_per_batch, num_samples, - 0.0, + cone_angle, scene_scale, ) diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index d5e49fe05..75cfd8f65 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -239,3 +239,9 @@ class NerfWModel(Model): if "mask" in batch: mask = batch["mask"].repeat(1, 1, 3) writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step) + + # this doesn't really make sense but do it anyway + fine_psnr = self.psnr(image, rgb_fine) + return { + "psnr": float(fine_psnr.item()), + } diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index bac7089c2..1d4fcb9d2 100644 --- a/nerfactory/models/tensorf.py +++ b/nerfactory/models/tensorf.py @@ -226,4 +226,6 @@ class TensoRFModel(Model): writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(psnr), step=step) - return psnr.item() + return { + "psnr": float(psnr.item()), + } diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index f261274a3..b62f1ab7c 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -222,4 +222,4 @@ class NeRFModel(Model): writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) - return fine_psnr.item() + return {"psnr": float(fine_psnr.item())} diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 85790cdcd..96c87e8ab 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -138,10 +138,10 @@ class Pipeline(nn.Module): image_idx = int(camera_ray_bundle.camera_indices[0, 0]) if self.world_size > 1: outputs = self.model.module.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs) + metrics_dict = self.model.module.log_test_image_outputs(image_idx, step, batch, outputs) else: outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = self.model.log_test_image_outputs(image_idx, step, batch, outputs) + metrics_dict = self.model.log_test_image_outputs(image_idx, step, batch, outputs) if step is not None: writer.put_time( name=EventName.TEST_RAYS_PER_SEC, @@ -149,8 +149,8 @@ class Pipeline(nn.Module): step=step, avg_over_steps=True, ) - # TODO(ethan): this function should probably return something? self.train() + return metrics_dict @abstractmethod def log_test_image_outputs(self) -> None: diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index 79d17dd48..86d445564 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -22,7 +22,7 @@ import logging from abc import abstractmethod from pathlib import Path from time import time -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional import imageio import numpy as np @@ -61,6 +61,7 @@ class EventType(enum.Enum): IMAGE = "write_image" SCALAR = "write_scalar" DICT = "write_scalar_dict" + CONFIG = "write_config" @check_main_thread @@ -93,7 +94,7 @@ def put_scalar(name: str, scalar: float, step: int): @check_main_thread -def put_dict(name: str, scalar_dict: float, step: int): +def put_dict(name: str, scalar_dict: Dict[str, Any], step: int): """Setter function to place a dictionary of scalars into the queue to be written out Args: @@ -104,6 +105,18 @@ def put_dict(name: str, scalar_dict: float, step: int): EVENT_STORAGE.append({"name": name, "write_type": EventType.DICT, "event": scalar_dict, "step": step}) +@check_main_thread +def put_config(name: str, config_dict: Dict[str, Any], step: int): + """Setter function to place a dictionary of scalars into the queue to be written out + + Args: + name: name of scalar dictionary + scalar_dict: values to write out + step: step associated with dict + """ + EVENT_STORAGE.append({"name": name, "write_type": EventType.CONFIG, "event": config_dict, "step": step}) + + @check_main_thread def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, update_eta: bool = False): """Setter function to place a time element into the queue to be written out. @@ -123,7 +136,6 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, GLOBAL_BUFFER["step"] = step curr_event = GLOBAL_BUFFER["events"].get(name, {"buffer": [], "avg": 0}) curr_buffer = curr_event["buffer"] - curr_avg = curr_event["avg"] if len(curr_buffer) >= GLOBAL_BUFFER["max_buffer_size"]: curr_buffer.pop(0) curr_buffer.append(duration) @@ -147,14 +159,10 @@ def write_out_storage(): for writer in EVENT_WRITERS: for event in EVENT_STORAGE: write_func = getattr(writer, event["write_type"].value) - if event["write_type"] == EventType.DICT: - write_func(event["event"], event["step"]) - if isinstance(writer, LocalWriter): - continue - else: - write_func(event["name"], event["event"], event["step"]) - if isinstance(writer, LocalWriter): - break + write_func(event["name"], event["event"], event["step"]) + # NOTE(ethan): why is the below statement needed? + if isinstance(writer, LocalWriter): + break EVENT_STORAGE.clear() @@ -212,15 +220,15 @@ class Writer: raise NotImplementedError @check_main_thread - def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int) -> None: + def write_scalar_dict(self, name: str, scalar_dict: Dict[str, Any], step: int) -> None: """Function that writes out all scalars from a given dictionary to the logger Args: scalar_dict: dictionary containing all scalar values with key names and quantities step: the time step to log """ - for name, scalar in scalar_dict.items(): - self.write_scalar(name, scalar, step) + for key, scalar in scalar_dict.items(): + self.write_scalar(name + "/" + key, float(scalar), step) class TimeWriter: @@ -258,7 +266,7 @@ class WandbWriter(Writer): def __init__(self, config: cfg.WandbWriterConfig): super().__init__(config.log_dir) - wandb.init(dir=config.log_dir) + wandb.init(project="nerfactory-project", dir=config.log_dir, reinit=True) def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None: image = torch.permute(image, (2, 0, 1)) @@ -267,6 +275,16 @@ class WandbWriter(Writer): def write_scalar(self, name: str, scalar: float, step: int) -> None: wandb.log({name: scalar}, step=step) + def write_config(self, name: str, config_dict: Dict[str, Any], step: int): + # pylint: disable=unused-argument + # pylint: disable=no-self-use + """Function that writes out the config to wandb + + Args: + config: config dictionary to write out + """ + wandb.config.update(config_dict) + @decorate_all([check_main_thread]) class TensorboardWriter(Writer): @@ -283,6 +301,14 @@ class TensorboardWriter(Writer): def write_scalar(self, name: str, scalar: float, step: int) -> None: self.tb_writer.add_scalar(name, scalar, step) + def write_config(self, name: str, config_dict: Dict[str, Any], step: int): # pylint: disable=unused-argument + """Function that writes out the config to tensorboard + + Args: + config: config dictionary to write out + """ + self.tb_writer.add_text("config", str(config_dict)) + def _cursorup(x: int): """utility tool to move the cursor up on the terminal @@ -353,6 +379,14 @@ class LocalWriter(Writer): self._update_header(latest_map, new_key) self._print_stats(latest_map) + def write_config(self, name: str, config_dict: Dict[str, Any], step: int): + """Function that writes out the config to local + + Args: + config: config dictionary to write out + """ + # TODO: implement this + def _consolidate_events(self): latest_map = {} new_key = False diff --git a/pyproject.toml b/pyproject.toml index 6e0ce4a3b..469019caf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ dependencies = [ "torchvision>=0.13.0", "tqdm==4.64.0", "u-msgpack-python>=2.4.1", - "wandb==0.13.2" + "wandb==0.13.3" ] [project.urls] diff --git a/scripts/run_eval.py b/scripts/run_eval.py index c7e31b771..84e1fea7e 100755 --- a/scripts/run_eval.py +++ b/scripts/run_eval.py @@ -86,7 +86,8 @@ def _render_stats_dict(pipeline: Pipeline) -> Dict[str, float]: with torch.no_grad(): image_idx = int(camera_ray_bundle.camera_indices[0, 0]) outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - psnr = pipeline.model.log_test_image_outputs(image_idx, step, batch, outputs) + metrics_dict = pipeline.model.log_test_image_outputs(image_idx, step, batch, outputs) + psnr = metrics_dict["psnr"] avg_rays_per_sec = _update_avg(avg_rays_per_sec, camera_ray_bundle.origins.shape[0] / t.duration, step) avg_psnr = _update_avg(avg_psnr, psnr, step) avg_fps = _update_avg(avg_fps, 1 / t.duration, step) diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 8926a753b..ac203ed43 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -30,6 +30,10 @@ def set_reduced_config(config: cfg.Config): data_directory=Path("tests/data/lego_test") ) + # use tensorboard logging instead of wandb + config.logging.writer[0].enable = True + config.logging.writer[1].enable = False + # reduce model factors config.pipeline.model.num_coarse_samples = 4 config.pipeline.model.num_importance_samples = 4