From 4350bc78d64306c0d1eb6389b0640517d86f51d3 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Fri, 20 May 2022 10:41:19 -0700 Subject: [PATCH 1/2] Remove density field --- mattport/nerf/fields/base.py | 32 ++++++++------------- mattport/nerf/fields/density_field.py | 40 --------------------------- mattport/nerf/graph/vanilla_nerf.py | 8 ++---- mattport/structures/rays.py | 30 +++++++++----------- 4 files changed, 28 insertions(+), 82 deletions(-) delete mode 100644 mattport/nerf/fields/density_field.py diff --git a/mattport/nerf/fields/base.py b/mattport/nerf/fields/base.py index e930b2174..cbea5dad8 100644 --- a/mattport/nerf/fields/base.py +++ b/mattport/nerf/fields/base.py @@ -3,21 +3,17 @@ Base class for the graphs. """ from abc import abstractmethod +from typing import Tuple from torch import nn +from torchtyping import TensorType -from mattport.structures.rays import PointSamples, RaySamples +from mattport.structures.rays import PointSamples class Field(nn.Module): """Base class for fields.""" - def __init__(self) -> None: - super().__init__() - # minimum weight for it to be valid. used to skip querying the field - # TODO(ethan): avoid the hardcoding - self.weights_threshold = 1e-4 - def density_fn(self, positions): """Returns only the density. Used primarily with the occupancy grid.""" point_samples = PointSamples(positions=positions) @@ -25,8 +21,15 @@ class Field(nn.Module): return density @abstractmethod - def get_density(self, point_samples: PointSamples): - """Computes and returns the densities.""" + def get_density(self, point_samples: PointSamples) -> Tuple[TensorType[..., 1], TensorType[..., "num_features"]]: + """Computes and returns the densities. + + Args: + point_samples (PointSamples): Samples locations to compute density. + + Returns: + Tuple[TensorType[...,1], TensorType[...,"num_features"]]: A tensor of densities and a tensor of features. + """ @abstractmethod def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): @@ -38,14 +41,3 @@ class Field(nn.Module): field_outputs = self.get_outputs(point_samples, density_embedding=density_embedding) field_outputs["density"] = density return field_outputs - - def forward_with_weight_pruning(self, ray_samples: RaySamples): - """Forward but with pruning based on weights, from the densities.""" - density, density_embedding = self.get_density(ray_samples) - weights = ray_samples.get_weights(density) - valid_mask_outputs = weights >= self.weights_threshold - field_outputs = self.get_outputs( - ray_samples, density_embedding=density_embedding, valid_mask=valid_mask_outputs - ) - field_outputs["density"] = density - return field_outputs diff --git a/mattport/nerf/fields/density_field.py b/mattport/nerf/fields/density_field.py deleted file mode 100644 index a31a0a67e..000000000 --- a/mattport/nerf/fields/density_field.py +++ /dev/null @@ -1,40 +0,0 @@ -""" -Density Field. -""" - -from torch import nn - -from torchtyping import TensorType -from mattport.nerf.field_modules.field_heads import DensityFieldHead -from mattport.structures.rays import RaySamples - - -class DensityField(nn.Module): - """Base class for fields.""" - - def __init__(self, model: nn.Sequential, density_activation: nn.Module = nn.Softplus()) -> None: - super().__init__() - self.model = model - self.density_head = DensityFieldHead(in_dim=self.model[-1].get_out_dim(), activation=density_activation) - - def process_point(self, in_tensor: TensorType): - """_summary_""" - features = self.model(in_tensor) - density = self.density_head(features) - return density, features - - def proccess_ray(self, rays: RaySamples): - """_summary_""" - densities, features = self.process_point(in_tensor=rays.positions) - - weights = rays.get_weights(densities) - return densities, weights, features - - def forward(self, in_tensor): - """summary""" - return self.process_point(in_tensor) - - def density_fn(self, positions): - """Returns only the density.""" - density, _ = self.process_point(in_tensor=positions) - return density diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index f1ba247ba..763d67747 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -15,7 +15,6 @@ from mattport.nerf.field_modules.encoding import NeRFEncoding from mattport.nerf.field_modules.field_heads import DensityFieldHead, RGBFieldHead from mattport.nerf.field_modules.mlp import MLP from mattport.nerf.fields.base import Field -from mattport.nerf.fields.density_field import DensityField from mattport.nerf.graph.base import Graph from mattport.nerf.loss import MSELoss from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer @@ -43,9 +42,6 @@ class NeRFField(Field): self.build_mlp_rgb() self.build_heads() - sequence = nn.Sequential(self.encoding_xyz, self.mlp_base) - self.density_field = DensityField(sequence) - def build_encodings(self): """Build the encodings.""" self.encoding_xyz = NeRFEncoding( @@ -95,7 +91,9 @@ class NeRFField(Field): if not valid_mask.any(): # empty mask return density, base_mlp_out - density[valid_mask], base_mlp_out[valid_mask] = self.density_field(positions[valid_mask]) + encoded_xyz = self.encoding_xyz(positions[valid_mask]) + base_mlp_out[valid_mask] = self.mlp_base(encoded_xyz) + density[valid_mask] = self.field_output_density(base_mlp_out[valid_mask]) return density, base_mlp_out def get_outputs(self, point_samples: PointSamples, density_embedding=None, valid_mask=None): diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 7593041ea..616c00850 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -15,22 +15,22 @@ from mattport.utils.misc import is_not_none class PointSamples: """Samples in space.""" - positions: TensorType[..., 3] = None - directions: TensorType[..., 3] = None - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[..., 1] = None + positions: TensorType[..., 3] = None # XYZ locations + directions: TensorType[..., 3] = None # Unit direction vector + camera_indices: TensorType[..., 1] = None # Camera index + valid_mask: TensorType[..., 1] = None # Rays that are valid @dataclass class RaySamples: """Samples along a ray""" - positions: TensorType[..., 3] = None - directions: TensorType[..., 3] = None - camera_indices: TensorType[..., 1] = None - valid_mask: TensorType[..., 1] = None - ts: TensorType[..., 1] = None - deltas: TensorType[..., 1] = None + positions: TensorType[..., 3] = None # XYZ locations + directions: TensorType[..., 3] = None # Unit direction vector + camera_indices: TensorType[..., 1] = None # Camera index + valid_mask: TensorType[..., 1] = None # Rays that are valid + ts: TensorType[..., 1] = None # "time steps", distances along ray + deltas: TensorType[..., 1] = None # "width" of each sample def to_point_samples(self) -> PointSamples: """Convert to PointSamples instance and return.""" @@ -67,14 +67,10 @@ class RaySamples: @dataclass class RayBundle: - """_summary_ + """A bundle of ray parameters.""" - Returns: - _type_: _description_ - """ - - origins: TensorType["num_rays", 3] - directions: TensorType["num_rays", 3] + origins: TensorType["num_rays", 3] # Ray origins + directions: TensorType["num_rays", 3] # camera_indices: Optional[TensorType["num_rays"]] = None nears: Optional[TensorType["num_rays"]] = None fars: Optional[TensorType["num_rays"]] = None From 2144fd152710e02dc3b15c35f570b1c8d826e78f Mon Sep 17 00:00:00 2001 From: evonneng Date: Fri, 20 May 2022 11:41:34 -0700 Subject: [PATCH 2/2] Enable stats tracker logging on full history (#91) * added logger ability to log everything without erasing --- configs/logging/default_logging.yaml | 5 +- mattport/nerf/graph/nerfw.py | 2 +- mattport/nerf/graph/vanilla_nerf.py | 2 +- mattport/nerf/trainer.py | 14 +++-- mattport/utils/writer.py | 85 +++++++++++++++++++--------- 5 files changed, 72 insertions(+), 36 deletions(-) diff --git a/configs/logging/default_logging.yaml b/configs/logging/default_logging.yaml index 8938d1233..9dfaf40e1 100644 --- a/configs/logging/default_logging.yaml +++ b/configs/logging/default_logging.yaml @@ -1,12 +1,13 @@ steps_per_log: 10 -max_history: 20 +max_buffer_size: 20 # writer logs losses and images per iteration writer: TensorboardWriter: save_dir: './' LocalWriter: save_dir: './' - stats_to_track: ["Data Load (time)", "Train Iter (time)", "Rays Per Sec (time)", "Test PSNR"] + stats_to_track: [ITER_LOAD_TIME, ITER_TRAIN_TIME, RAYS_PER_SEC, CURR_TEST_PSNR] + max_log_size: 20 # if 0, logs everything with no erasing # profiler logs run times of functions and prints at end of training enable_profiler: True \ No newline at end of file diff --git a/mattport/nerf/graph/nerfw.py b/mattport/nerf/graph/nerfw.py index b31061266..0b1a09b1a 100644 --- a/mattport/nerf/graph/nerfw.py +++ b/mattport/nerf/graph/nerfw.py @@ -212,4 +212,4 @@ class NerfWGraph(NeRFGraph): row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2) combined_image = torch.cat([row0, row1, row2], dim=-3) - writer.put_scalar(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") + writer.put_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img") diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 763d67747..2bb77c109 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -249,6 +249,6 @@ class NeRFGraph(Graph): writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim") writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips") - writer.put_scalar(name="Test PSNR", scalar=float(fine_psnr), step=step) + writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step) return fine_psnr.item() diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index e6f654a24..e86be776d 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -202,13 +202,17 @@ class Trainer: data_start = time() batch = next(iter_dataset) writer.put_time( - name="Data Load (time)", start_time=data_start, end_time=time(), step=step, avg_over_iters=True + name=writer.EventName.ITER_LOAD_TIME, + start_time=data_start, + end_time=time(), + step=step, + avg_over_iters=True, ) iter_start = time() loss_dict = self.train_iteration(batch, step) writer.put_time( - name="Rays Per Sec (time)", + name=writer.EventName.RAYS_PER_SEC, start_time=iter_start, end_time=time(), step=step, @@ -216,7 +220,7 @@ class Trainer: avg_over_batch=batch["indices"].shape[0], ) writer.put_time( - name="Train Iter (time)", + name=writer.EventName.ITER_TRAIN_TIME, start_time=iter_start, end_time=time(), step=step, @@ -234,7 +238,9 @@ class Trainer: _ = self.test_image(image_idx=image_idx, step=step) self._write_out_storage(step) - writer.put_time(name="Train Total (time)", start_time=train_start, end_time=time(), step=num_iterations) + writer.put_time( + name=writer.EventName.TOTAL_TRAIN_TIME, start_time=train_start, end_time=time(), step=num_iterations + ) self._write_out_storage(num_iterations) def _write_out_storage(self, step): diff --git a/mattport/utils/writer.py b/mattport/utils/writer.py index ecacd7567..47130761a 100644 --- a/mattport/utils/writer.py +++ b/mattport/utils/writer.py @@ -25,12 +25,20 @@ EVENT_STORAGE = [] GLOBAL_BUFFER = {} -class EventType(enum.Enum): - """Possible Event types that can be logged +class EventName(enum.Enum): + """Names of possible events that can be logged via Local Writer for convenience. + see config/logging/default_logging.yaml""" - Args: - enum (_type_): _description_ - """ + ITER_LOAD_TIME = "Data Load (time)" + ITER_TRAIN_TIME = "Train Iter (time)" + RAYS_PER_SEC = "Rays Per Sec (time)" + TOTAL_TRAIN_TIME = "Train Total (time)" + ETA = "ETA (time)" + CURR_TEST_PSNR = "Test PSNR" + + +class EventType(enum.Enum): + """Possible Event types and their associated write function""" IMAGE = "write_image" SCALAR = "write_scalar" @@ -40,6 +48,9 @@ class EventType(enum.Enum): @check_main_thread def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None): """Setter function to place images into the queue to be written out""" + if isinstance(name, EventName): + name = name.value + EVENT_STORAGE.append( {"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix} ) @@ -48,6 +59,9 @@ def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = No @check_main_thread def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None): """Setter function to place scalars into the queue to be written out""" + if isinstance(name, EventName): + name = name.value + GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] GLOBAL_BUFFER["events"][name] = scalar EVENT_STORAGE.append( @@ -89,6 +103,9 @@ def put_time( avg_over_batch (int): if set, the size of the batch for which we take the average over (batch/second) update_eta (bool): if True, update the ETA. should only be set for the training iterations/s """ + if isinstance(name, EventName): + name = name.value + GLOBAL_BUFFER["step"] = step GLOBAL_BUFFER["new_key"] = not name in GLOBAL_BUFFER["events"] or GLOBAL_BUFFER["new_key"] val = end_time - start_time @@ -99,7 +116,7 @@ def put_time( 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_history"]: + if len(curr_buffer) >= GLOBAL_BUFFER["max_buffer_size"]: curr_buffer.pop(0) curr_buffer.append(val) curr_avg = sum(curr_buffer) / len(curr_buffer) @@ -114,7 +131,7 @@ def put_time( remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] GLOBAL_BUFFER["events"]["ETA"] = remain_time - put_scalar("ETA", remain_time, step, group, prefix) + put_scalar(EventName.ETA, remain_time, step, group, prefix) @check_main_thread @@ -140,16 +157,16 @@ def setup_event_writers(config: DictConfig) -> None: logging_configs = config.logging.writer for writer_type in logging_configs: writer_class = getattr(mattport.utils.writer, writer_type) - curr_config = logging_configs[writer_type] + writer_config = logging_configs[writer_type] if writer_type == "LocalWriter": - curr_writer = writer_class(curr_config.save_dir, curr_config.stats_to_track, config) + curr_writer = writer_class(writer_config.save_dir, writer_config.stats_to_track, writer_config.max_log_size) else: - curr_writer = writer_class(curr_config.save_dir) + curr_writer = writer_class(writer_config.save_dir) EVENT_WRITERS.append(curr_writer) ## configure all the global buffer basic information GLOBAL_BUFFER["max_iter"] = config.graph.max_num_iterations - GLOBAL_BUFFER["max_history"] = config.logging.max_history + GLOBAL_BUFFER["max_buffer_size"] = config.logging.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.logging.steps_per_log GLOBAL_BUFFER["new_key"] = True GLOBAL_BUFFER["events"] = {} @@ -278,12 +295,17 @@ def _consolidate_events(): class LocalWriter(Writer): """Local Writer Class""" - def __init__(self, save_dir: str, stats_to_track: ListConfig, config: DictConfig): + def __init__(self, save_dir: str, stats_to_track: ListConfig, max_log_size: int = 0): + """ + Args: + stats_to_track (ListConfig): the names of stats that should be logged. + max_log size (int): max number of lines that will be logged to teminal. + """ super().__init__(save_dir) - self.stats_to_track = stats_to_track + self.stats_to_track = [EventName[name].value for name in stats_to_track] + self.max_log_size = max_log_size self.past_stats = [] self.max_mssg_len = 0 - self.config = config def write_image( self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None @@ -301,18 +323,22 @@ class LocalWriter(Writer): def handle_header(self, latest_map): """helper to handle the printing of the header labels""" - if len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]: + if (not self.max_log_size and self.max_mssg_len == 0) or ( + self.max_log_size and (len(self.past_stats) == 0 or GLOBAL_BUFFER["new_key"]) + ): mssg = f"{'Step (% Done)':<20}" for name, _ in latest_map.items(): if name in self.stats_to_track: mssg += f"{name:<20} " - if len(self.past_stats) > 1: - _cursorup(len(self.past_stats) - 1) + if self.max_log_size: + if len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) print(mssg) print("-" * len(mssg)) - if len(self.past_stats) > 0: - for mssg in self.past_stats: - print(mssg) + if self.max_log_size: + if len(self.past_stats) > 0: + for mssg in self.past_stats: + print(mssg) self.max_mssg_len = len(mssg) def handle_stats(self, latest_map, padding=" "): @@ -330,12 +356,15 @@ class LocalWriter(Writer): curr_mssg += f"{v:<20} " # update the history buffer - if len(self.past_stats) > GLOBAL_BUFFER["max_history"]: - self.past_stats.pop(0) - _cursorup(len(self.past_stats)) - elif len(self.past_stats) > 1: - _cursorup(len(self.past_stats) - 1) + if self.max_log_size: + if len(self.past_stats) > self.max_log_size: + self.past_stats.pop(0) + _cursorup(len(self.past_stats)) + elif len(self.past_stats) > 1: + _cursorup(len(self.past_stats) - 1) - for mssg in self.past_stats: - print(f"{mssg:{padding}<{self.max_mssg_len}}") - self.past_stats.append(curr_mssg) + for mssg in self.past_stats: + print(f"{mssg:{padding}<{self.max_mssg_len}}") + self.past_stats.append(curr_mssg) + else: + print(curr_mssg)