From 3226a1c989dfe792a99e0eb505f9dd45b5986d49 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Thu, 26 May 2022 10:18:08 -0400 Subject: [PATCH] CPU usage improvement and minor changes (#108) * cpu fixes, masking logic but commented out, nerfw and semanticnerf in progress * pylint --- .../dataset/friends_TBBT-big_living_room.yaml | 2 +- configs/data/default_data.yaml | 2 +- configs/graph/default_graph.yaml | 8 ---- configs/instant_ngp_tcnn.yaml | 10 +++-- mattport/nerf/dataset/format/friends.py | 44 ++++++++++++++----- mattport/nerf/dataset/image_dataset.py | 7 +++ mattport/nerf/fields/instant_ngp_field.py | 3 +- mattport/nerf/graph/base.py | 10 +++-- mattport/nerf/graph/instant_ngp.py | 2 +- mattport/nerf/graph/nerfw.py | 30 ++++++------- mattport/nerf/graph/semantic_nerf.py | 4 +- mattport/nerf/graph/vanilla_nerf.py | 4 +- mattport/nerf/optimizers.py | 16 +++++-- mattport/nerf/ray_generator.py | 17 ++++--- mattport/nerf/trainer.py | 15 ++++--- mattport/structures/cameras.py | 2 +- mattport/structures/rays.py | 7 +++ .../nerf/dataloader/configs/blender_lego.yaml | 2 +- 18 files changed, 118 insertions(+), 67 deletions(-) diff --git a/configs/data/dataset/friends_TBBT-big_living_room.yaml b/configs/data/dataset/friends_TBBT-big_living_room.yaml index 539beca58..07981cf10 100644 --- a/configs/data/dataset/friends_TBBT-big_living_room.yaml +++ b/configs/data/dataset/friends_TBBT-big_living_room.yaml @@ -1,3 +1,3 @@ data_directory: data/friends/TBBT-big_living_room dataset_format: friends -downscale_factor: 8 +downscale_factor: 1 diff --git a/configs/data/default_data.yaml b/configs/data/default_data.yaml index a547545c5..468770169 100644 --- a/configs/data/default_data.yaml +++ b/configs/data/default_data.yaml @@ -9,7 +9,7 @@ val_num_rays_per_chunk: 4096 # dataloader options dataloader: - num_workers: 4 + num_workers: 1 shuffle: true num_images_to_sample_from: 8 num_rays_per_batch: 1024 diff --git a/configs/graph/default_graph.yaml b/configs/graph/default_graph.yaml index ee08bc912..d0010cc7c 100644 --- a/configs/graph/default_graph.yaml +++ b/configs/graph/default_graph.yaml @@ -30,11 +30,3 @@ param_groups: _target_: mattport.nerf.optimizers.ExponentialDecaySchedule lr_final: 0.000005 max_steps: 1000000 - - cameras: - optimizer: - _target_: torch.optim.SGD - lr: 1e-4 - scheduler: - _target_: torch.optim.lr_scheduler.ExponentialLR - gamma: 0.1 diff --git a/configs/instant_ngp_tcnn.yaml b/configs/instant_ngp_tcnn.yaml index e5778354c..c699d341d 100644 --- a/configs/instant_ngp_tcnn.yaml +++ b/configs/instant_ngp_tcnn.yaml @@ -3,6 +3,10 @@ defaults: - override graph: instant_ngp_graph - _self_ method_name: instant_ngp -data.dataloader.num_rays_per_batch: 4096 -data.val_num_rays_per_chunk: 16384 -graph.network.field_implementation: tcnn \ No newline at end of file +data: + dataloader: + num_rays_per_batch: 4096 + val_num_rays_per_chunk: 16384 +graph: + network: + field_implementation: tcnn diff --git a/mattport/nerf/dataset/format/friends.py b/mattport/nerf/dataset/format/friends.py index 552bdb858..f8c52de67 100644 --- a/mattport/nerf/dataset/format/friends.py +++ b/mattport/nerf/dataset/format/friends.py @@ -8,11 +8,33 @@ import numpy as np import torch from mattport.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary -from mattport.nerf.dataset.structs import DatasetInputs, PointCloud, Semantics +from mattport.nerf.dataset.structs import DatasetInputs, PointCloud, SceneBounds, Semantics from mattport.utils import profiler from mattport.utils.io import load_from_json +def get_aabb_and_transform(basedir): + """Returns the aabb and pointcloud transform from the threejs.json file.""" + filename = os.path.join(basedir, "threejs.json") + assert os.path.exists(filename) + data = load_from_json(filename) + + # point cloud transformation + transposed_point_cloud_transform = np.array(data["object"]["children"][0]["matrix"]).reshape(4, 4).T + assert transposed_point_cloud_transform[3, 3] == 1.0 + + # bbox transformation + bbox_transform = np.array(data["object"]["children"][1]["matrix"]).reshape(4, 4).T + w, h, d = data["geometries"][1]["width"], data["geometries"][1]["height"], data["geometries"][1]["depth"] + temp = np.array([w, h, d]) / 2.0 + bbox = np.array([-temp, temp]) + bbox = np.concatenate([bbox, np.ones_like(bbox[:, 0:1])], axis=1) + bbox = (bbox_transform @ bbox.T).T[:, 0:3] + + aabb = bbox # rename to aabb because it's an axis-aligned bounding box + return torch.from_numpy(aabb).float(), torch.from_numpy(transposed_point_cloud_transform).float() + + @profiler.time_function def load_friends_data(basedir, downscale_factor=1, split="train", include_semantics=True, include_point_cloud=False): """_summary_ @@ -38,6 +60,10 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant # TODO: handle the splits differently image_filenames = [os.path.join(basedir, "images", image_path) for image_path in image_paths] + # -- set the bounding box --- + aabb, transposed_point_cloud_transform = get_aabb_and_transform(basedir) + scene_bounds = SceneBounds(aabb=aabb) + # --- intrinsics --- cameras_data = read_cameras_binary(os.path.join(basedir, "colmap", "cameras.bin")) intrinsics = [] @@ -61,8 +87,9 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant camera_to_world.append(c2w) camera_to_world = torch.tensor(np.array(camera_to_world)).float() camera_to_world = torch.inverse(camera_to_world) - camera_to_world = camera_to_world[:, :3] camera_to_world[..., 1:3] *= -1 + camera_to_world = transposed_point_cloud_transform @ camera_to_world + camera_to_world = camera_to_world[:, :3] # --- masks to mask out things (e.g., people) --- @@ -87,19 +114,15 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant thing_filenames=thing_filenames, ) - # --- possibly transform --- - # TODO: - - # -- set the bounding box --- - # TODO: - # Possibly include the sparse point cloud from COLMAP in the dataset inputs. # NOTE(ethan): this will be common across the different splits. point_cloud = PointCloud() if include_point_cloud: points_3d = read_pointsTD_binary(os.path.join(basedir, "colmap", "points3D.bin")) - xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])) - rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])) + xyz = torch.tensor(np.array([p_value.xyz for p_id, p_value in points_3d.items()])).float() + rgb = torch.tensor(np.array([p_value.rgb for p_id, p_value in points_3d.items()])).float() + xyz_h = torch.cat([xyz, torch.ones_like(xyz[..., :1])], -1) + xyz = (xyz_h @ transposed_point_cloud_transform.T)[..., :3] point_cloud.xyz = xyz point_cloud.rgb = rgb @@ -110,5 +133,6 @@ def load_friends_data(basedir, downscale_factor=1, split="train", include_semant camera_to_world=camera_to_world, semantics=semantics, point_cloud=point_cloud, + scene_bounds=scene_bounds, ) return dataset_inputs diff --git a/mattport/nerf/dataset/image_dataset.py b/mattport/nerf/dataset/image_dataset.py index 907d45e7b..d0c19e2af 100644 --- a/mattport/nerf/dataset/image_dataset.py +++ b/mattport/nerf/dataset/image_dataset.py @@ -27,6 +27,13 @@ def collate_batch(batch_list, num_rays_per_batch, keep_full_image: bool = False) batch = default_collate(batch_list) num_images, image_height, image_width, _ = batch["image"].shape + + # # checking where mask is not zero + # # TODO(ethan): add this to the code with a config + # nonzero_indices = torch.nonzero(batch["mask"], as_tuple=False) + # chosen_indices = random.sample(range(len(nonzero_indices)), k=num_rays_per_batch) + # indices = nonzero_indices[chosen_indices] + indices = torch.floor( torch.rand((num_rays_per_batch, 3)) * torch.tensor([num_images, image_height, image_width]) ).long() diff --git a/mattport/nerf/fields/instant_ngp_field.py b/mattport/nerf/fields/instant_ngp_field.py index 092635136..d50c13ea9 100644 --- a/mattport/nerf/fields/instant_ngp_field.py +++ b/mattport/nerf/fields/instant_ngp_field.py @@ -10,6 +10,7 @@ import torch.nn.functional as F from torch.nn.parameter import Parameter from mattport.nerf.field_modules.encoding import Encoding, HashEncoding, SHEncoding +from mattport.nerf.field_modules.field_heads import FieldHeadNames from mattport.nerf.fields.base import Field from mattport.nerf.fields.nerf_field import NeRFField from mattport.structures.rays import PointSamples @@ -112,7 +113,7 @@ class TCNNInstantNGPField(Field): h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) h = self.mlp_head(h).view(*point_samples.directions.shape[:-1], -1).to(dtype) rgb = torch.sigmoid(h) - return {"rgb": rgb} + return {FieldHeadNames.RGB: rgb} class TorchInstantNGPField(NeRFField): diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py index fa0ee1cec..72dfe9b20 100644 --- a/mattport/nerf/graph/base.py +++ b/mattport/nerf/graph/base.py @@ -107,7 +107,7 @@ class Graph(AbstractGraph): # during training, keep only the rays that intersect the scene. discard the rest valid_mask = intersected_ray_bundle.valid_mask masked_intersected_ray_bundle = intersected_ray_bundle.get_masked_ray_bundle(valid_mask) - masked_batch = get_masked_dict(batch, valid_mask) + masked_batch = get_masked_dict(batch, valid_mask) # NOTE(ethan): this is really slow if on CPU! outputs = self.get_outputs(masked_intersected_ray_bundle) loss_dict = self.get_loss_dict(outputs=outputs, batch=masked_batch) return outputs, loss_dict @@ -128,6 +128,9 @@ class Graph(AbstractGraph): @torch.no_grad() def get_outputs_for_camera(self, intrinsics, camera_to_world, chunk_size=1024, training_camera_index=0): """Takes in camera parameters and computes the output of the graph.""" + # NOTE(ethan): this function has a spike in CPU usage + + device = self.get_device() assert len(intrinsics.shape) == 1 num_intrinsics_params = len(intrinsics) camera_class = get_camera_model(num_intrinsics_params) @@ -135,9 +138,9 @@ class Graph(AbstractGraph): camera_ray_bundle = camera.generate_camera_rays() # TODO(ethan): decide how to properly handle the image indices for validation images camera_ray_bundle.set_camera_indices(camera_index=training_camera_index) + camera_ray_bundle.move_to_device(device) image_height, image_width = camera_ray_bundle.origins.shape[:2] - device = self.get_device() num_rays = len(camera_ray_bundle) outputs = {} @@ -146,7 +149,6 @@ class Graph(AbstractGraph): start_idx = i end_idx = i + chunk_size original_ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) - original_ray_bundle.move_to_device(device) intersected_ray_bundle = self.collider(original_ray_bundle) outputs = self.get_outputs(intersected_ray_bundle) for output_name, output in outputs.items(): @@ -156,5 +158,5 @@ class Graph(AbstractGraph): return outputs @abstractmethod - def log_test_image_outputs(self, image_idx, step, image, outputs): + def log_test_image_outputs(self, image_idx, step, image, mask, outputs): """Writes the test image outputs.""" diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py index 783f06de1..ad67c9823 100644 --- a/mattport/nerf/graph/instant_ngp.py +++ b/mattport/nerf/graph/instant_ngp.py @@ -118,7 +118,7 @@ class NGPGraph(Graph): loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, outputs): + def log_test_image_outputs(self, image_idx, step, image, mask, outputs): rgb = outputs["rgb"] acc = visualization.apply_colormap(outputs["accumulation"]) depth = visualization.apply_depth_colormap( diff --git a/mattport/nerf/graph/nerfw.py b/mattport/nerf/graph/nerfw.py index 0b1a09b1a..b325df13a 100644 --- a/mattport/nerf/graph/nerfw.py +++ b/mattport/nerf/graph/nerfw.py @@ -22,28 +22,27 @@ from mattport.structures.rays import RayBundle, RaySamples from mattport.utils import visualization, writer -class NerfWField(NeRFField): +class NerfWField(Field): """The NeRF-W field which has appearance and transient conditioning.""" def __init__( self, - num_layers=8, - layer_width=256, - skip_connections: Tuple = (4,), - num_images: int = 0, - appearance_embedding_dim: int = 48, - transient_embedding_dim: int = 16, + position_encoding: Encoding = Identity(in_dim=3), + direction_encoding: Encoding = Identity(in_dim=3), + base_mlp_num_layers: int = 8, + base_mlp_layer_width: int = 256, + head_mlp_num_layers: int = 2, + head_mlp_layer_width: int = 128, + skip_connections: Tuple[int] = (4,), + field_heads: Tuple[FieldHead] = (RGBFieldHead(),), ) -> None: - assert num_images > 0 + super().__init__() self.num_images = num_images - self.appearance_embedding_dim = appearance_embedding_dim - self.transient_embedding_dim = transient_embedding_dim - super().__init__(num_layers, layer_width, skip_connections) + self.appearance_embedding_dim = 48 + self.transient_embedding_dim = 16 - def build_mlp_base(self): self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim) self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim) - super().build_mlp_base() self.mlp_transient = MLP( in_dim=self.mlp_base.get_out_dim() + self.embedding_transient.get_out_dim(), out_dim=self.layer_width // 2, @@ -52,8 +51,7 @@ class NerfWField(NeRFField): activation=nn.ReLU(), ) - def build_mlp_rgb(self): - self.mlp_rgb = MLP( + self.mlp_head = MLP( in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim() + self.embedding_appearance.get_out_dim(), @@ -193,7 +191,7 @@ class NerfWGraph(NeRFGraph): loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, outputs): + def log_test_image_outputs(self, image_idx, step, image, mask, outputs): rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] rgb_fine_static = outputs["rgb_fine_static"] diff --git a/mattport/nerf/graph/semantic_nerf.py b/mattport/nerf/graph/semantic_nerf.py index 726c54d77..494ec6d1d 100644 --- a/mattport/nerf/graph/semantic_nerf.py +++ b/mattport/nerf/graph/semantic_nerf.py @@ -149,8 +149,8 @@ class SemanticNerfGraph(NeRFGraph): loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, outputs): - super().log_test_image_outputs(image_idx, step, image, outputs) + def log_test_image_outputs(self, image_idx, step, image, mask, outputs): + super().log_test_image_outputs(image_idx, step, image, mask, outputs) semantic_labels = torch.argmax(torch.nn.functional.softmax(outputs["semantic_logits"], dim=-1), dim=-1) / len( self.stuff_classes ) diff --git a/mattport/nerf/graph/vanilla_nerf.py b/mattport/nerf/graph/vanilla_nerf.py index 875f6c80a..860603b0f 100644 --- a/mattport/nerf/graph/vanilla_nerf.py +++ b/mattport/nerf/graph/vanilla_nerf.py @@ -129,7 +129,7 @@ class NeRFGraph(Graph): loss_dict["aggregated_loss"] = self.get_aggregated_loss_from_loss_dict(loss_dict) return loss_dict - def log_test_image_outputs(self, image_idx, step, image, outputs): + def log_test_image_outputs(self, image_idx, step, image, mask, outputs): rgb_coarse = outputs["rgb_coarse"] rgb_fine = outputs["rgb_fine"] acc_coarse = visualization.apply_colormap(outputs["accumulation_coarse"]) @@ -150,10 +150,12 @@ class NeRFGraph(Graph): combined_rgb = torch.cat([image, rgb_coarse, rgb_fine], dim=1) combined_acc = torch.cat([acc_coarse, acc_fine], dim=1) combined_depth = torch.cat([depth_coarse, depth_fine], dim=1) + mask = visualization.apply_depth_colormap(mask[..., None]) writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img") writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation") writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth") + writer.put_image(name=f"image_idx_{image_idx}", image=mask, step=step, group="mask") # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] diff --git a/mattport/nerf/optimizers.py b/mattport/nerf/optimizers.py index 181ac1517..539db4675 100644 --- a/mattport/nerf/optimizers.py +++ b/mattport/nerf/optimizers.py @@ -14,16 +14,26 @@ from mattport.utils import writer class ExponentialDecaySchedule(LambdaLR): - """Exponential learning rate decay function.""" + """Exponential learning rate decay function. + See https://github.com/google-research/google-research/blob/ + fd2cea8cdd86b3ed2c640cbe5561707639e682f3/jaxnerf/nerf/utils.py#L360 + for details. + """ - def __init__(self, optimizer, lr_init, lr_final, max_steps) -> None: + def __init__(self, optimizer, lr_init, lr_final, max_steps, lr_delay_steps=0, lr_delay_mult=1.0) -> None: def func(step): + if lr_delay_steps > 0: + delay_rate = lr_delay_mult + (1 - lr_delay_mult) * np.sin( + 0.5 * np.pi * np.clip(step / lr_delay_steps, 0, 1) + ) + else: + delay_rate = 1.0 t = np.clip(step / max_steps, 0, 1) log_lerp = np.exp(np.log(lr_init) * (1 - t) + np.log(lr_final) * t) multiplier = ( log_lerp / lr_init ) # divided by lr_init because the multiplier is with the initial learning rate - return multiplier + return delay_rate * multiplier super().__init__(optimizer, lr_lambda=func) diff --git a/mattport/nerf/ray_generator.py b/mattport/nerf/ray_generator.py index 50c74910f..fe4d97098 100644 --- a/mattport/nerf/ray_generator.py +++ b/mattport/nerf/ray_generator.py @@ -30,6 +30,12 @@ class RayGenerator(nn.Module): self.camera_to_world = nn.Parameter(camera_to_world, requires_grad=False) # TODO(ethan): add learnable parameters that are deltas on the intrinsics and camera_to_world parameters + # NOTE(ethan): we currently assume all images have the same height and width + camera_index = 0 + self.camera_class = get_camera_model(self.num_intrinsics_params) + camera = self.camera_class(*self.intrinsics[camera_index].tolist()) + self.image_coords = nn.Parameter(camera.get_image_coords(), requires_grad=False) + def forward(self, ray_indices: TensorType["num_rays", 3]) -> RayBundle: """Index into the cameras to generate the rays. @@ -41,13 +47,10 @@ class RayGenerator(nn.Module): x = ray_indices[:, 2] # col indices intrinsics = self.intrinsics[c] camera_to_world = self.camera_to_world[c] - # NOTE(ethan): we currently assume all images have the same height and width - camera_index = 0 - camera_class = get_camera_model(self.num_intrinsics_params) - camera = camera_class(*self.intrinsics[camera_index].tolist()) - image_coords = camera.get_image_coords().to(ray_indices.device) # NOTE(ethan): I can actually avoid this step - coords = image_coords[y, x] + coords = self.image_coords[y, x] - ray_bundle = camera_class.generate_rays(intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords) + ray_bundle = self.camera_class.generate_rays( + intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords + ) ray_bundle.camera_indices = c return ray_bundle diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py index b77d14a41..b820bf9dd 100644 --- a/mattport/nerf/trainer.py +++ b/mattport/nerf/trainer.py @@ -21,6 +21,7 @@ from mattport.nerf.optimizers import Optimizers from mattport.utils import profiler, writer from mattport.utils.callbacks import update_occupancy from mattport.utils.decorators import check_main_thread +from mattport.utils.misc import get_dict_to_torch logging.getLogger("PIL").setLevel(logging.WARNING) @@ -194,14 +195,13 @@ class Trainer: def train(self) -> None: """_summary_""" - # self.graph.train() - # assert self.graph.training, "Graph is not in train mode!" train_start = time() num_iterations = self.config.graph.max_num_iterations iter_dataset = iter(self.train_dataloader) for step in range(self.start_step, self.start_step + num_iterations): data_start = time() batch = next(iter_dataset) + batch = get_dict_to_torch(batch, device=self.device) writer.put_time( name=writer.EventName.ITER_LOAD_TIME, start_time=data_start, @@ -257,9 +257,8 @@ class Trainer: @profiler.time_function def train_iteration(self, batch: dict, step: int, _callback: List[Callable] = None): """Run one iteration with a batch of inputs.""" - # move batch to correct device - ray_indices = batch["indices"].to(self.device) - _, loss_dict = self.graph(ray_indices, batch=batch, step=step) + ray_indices = batch["indices"] + _, loss_dict = self.graph.forward(ray_indices, batch=batch, step=step) loss = loss_dict["aggregated_loss"] self.optimizers.zero_grad_all() loss.backward() @@ -281,7 +280,9 @@ class Trainer: outputs = self.graph.get_outputs_for_camera( intrinsics, camera_to_world, chunk_size=chunk_size, training_camera_index=training_camera_index ) - image = self.val_image_dataset[image_idx]["image"].to(self.device) - psnr = self.graph.log_test_image_outputs(image_idx, step, image, outputs) + val_image_data = self.val_image_dataset[image_idx] + image = val_image_data["image"].to(self.device) + mask = val_image_data["mask"].to(self.device) + psnr = self.graph.log_test_image_outputs(image_idx, step, image, mask, outputs) self.graph.train() return psnr diff --git a/mattport/structures/cameras.py b/mattport/structures/cameras.py index 60c721246..7d8c43975 100644 --- a/mattport/structures/cameras.py +++ b/mattport/structures/cameras.py @@ -106,7 +106,7 @@ class Camera: """ height = self.get_image_height() width = self.get_image_width() - intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, height, 1) # (num_rays, num_intrinsics_params) + intrinsics = self.get_intrinsics().unsqueeze(0).repeat(height, width, 1) # (num_rays, num_intrinsics_params) camera_to_world = self.camera_to_world.unsqueeze(0).repeat(height, width, 1, 1) # (num_rays, 3, 4) coords = self.get_image_coords() ray_bundle = self.generate_rays(intrinsics, camera_to_world, coords) diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py index 9b2589b2f..62baba364 100644 --- a/mattport/structures/rays.py +++ b/mattport/structures/rays.py @@ -89,12 +89,19 @@ class RaySamples: delta_density = self.deltas * densities[..., 0] alphas = 1 - torch.exp(-delta_density) + # mip-nerf version of transmittance calculation: transmittance = torch.cumsum(delta_density[..., :-1], dim=-1) transmittance = torch.cat( [torch.zeros((*transmittance.shape[:1], 1)).to(densities.device), transmittance], axis=-1 ) transmittance = torch.exp(-transmittance) # [..., "num_samples"] + # most nerf codebases do the following: + # transmittance = torch.cat( + # [torch.ones((*alphas.shape[:1], 1)).to(densities.device), 1.0 - alphas + 1e-10], dim=-1 + # ) + # transmittance = torch.cumprod(transmittance, dim=-1)[..., :-1] # [..., "num_samples"] + weights = alphas * transmittance # [..., "num_samples"] return weights diff --git a/tests/nerf/dataloader/configs/blender_lego.yaml b/tests/nerf/dataloader/configs/blender_lego.yaml index d6dae7b05..5fb9a3164 100644 --- a/tests/nerf/dataloader/configs/blender_lego.yaml +++ b/tests/nerf/dataloader/configs/blender_lego.yaml @@ -6,7 +6,7 @@ dataset: # dataloader options dataloader: - num_workers: 0 + num_workers: 1 shuffle: true num_images_to_sample_from: 8 num_rays_per_batch: 1024