From a10eb5caef51efc3c3c042752bd2122a7e279881 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 2 May 2022 13:27:30 -0700 Subject: [PATCH 1/2] Add basic NGP encoder --- .vscode/launch.json | 15 ++- configs/graph/instant_ngp_graph.yaml | 29 +++++ configs/instant_ngp.yaml | 23 ++++ mattport/nerf/field_modules/encoding.py | 94 +++++++++++++- mattport/nerf/graph/instant_ngp.py | 161 +++++++++++++++++++++++- 5 files changed, 313 insertions(+), 9 deletions(-) create mode 100644 configs/graph/instant_ngp_graph.yaml create mode 100644 configs/instant_ngp.yaml diff --git a/.vscode/launch.json b/.vscode/launch.json index 980221899..e4e2c9c40 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -3,7 +3,8 @@ // Hover to view descriptions of existing attributes. // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 "version": "0.2.0", - "configurations": [{ + "configurations": [ + { "name": "Python: Current File", "type": "python", "request": "launch", @@ -27,6 +28,16 @@ "args": [ "machine_config.num_gpus=0" ] - } + }, + { + "name": "Python: run_train ngp (gpu)", + "type": "python", + "request": "launch", + "program": "scripts/run_train.py", + "console": "integratedTerminal", + "args": [ + "--config-name=instant_ngp.yaml" + ] + }, ] } \ No newline at end of file diff --git a/configs/graph/instant_ngp_graph.yaml b/configs/graph/instant_ngp_graph.yaml new file mode 100644 index 000000000..49a0c188f --- /dev/null +++ b/configs/graph/instant_ngp_graph.yaml @@ -0,0 +1,29 @@ +# saving model specifications +model_dir: /tmp/mattport_models/ +steps_per_save: null +steps_per_test: 100 +max_num_iterations: 1000000 + +# pretrained loading options +resume_train: + load_dir: # pretrained model directory + load_step: # pretrained model's step number + +# NOTE(ethan): currently setting to nothing b/c we're still designing the config format... +network: + _target_: mattport.nerf.graph.instant_ngp.NGPGraph + near_plane: 1.0 + far_plane: 6.0 + num_coarse_samples: 64 + num_importance_samples: 64 + +# optimizer options for network +param_groups: + fields: + optimizer: + type: Adam + lr: 1e-2 + scheduler: + type: StepLR + step_size: 30 # Period of learning rate decay. + gamma: 0.1 # Multiplicative factor of learning rate decay. Default: 0.1. \ No newline at end of file diff --git a/configs/instant_ngp.yaml b/configs/instant_ngp.yaml new file mode 100644 index 000000000..cae7eca18 --- /dev/null +++ b/configs/instant_ngp.yaml @@ -0,0 +1,23 @@ +defaults: + # This makes the logging module logs colorful + - data: default_data.yaml + - graph: instant_ngp_graph.yaml + - logging: default_logging.yaml + - override hydra/job_logging: colorlog + - override hydra/hydra_logging: colorlog + +hydra: + run: + # Configure output dir of each experiment programmatically from the arguments + # Example "outputs/mnist/classifier/baseline/2021-03-10-141516" + dir: outputs/${experiment_name}/${method_name}/${now:%Y-%m-%d_%H%M%S} + +experiment_name: blender_lego +method_name: instant_ngp + +# machine options +machine_config: + num_gpus: 1 + num_machines: 1 + machine_rank: 0 + dist_url: auto diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index 0ff7bcb26..4e9f2ba85 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -1,6 +1,7 @@ """ Encoding functions """ +import numpy as np import torch from torchtyping import TensorType from mattport.nerf.field_modules.base import FieldModule @@ -141,13 +142,96 @@ class RFFEncoding(Encoding): class HashEncoding(Encoding): """Hash encoding""" - def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - """Encodes an input tensor. + def __init__( + self, + num_levels: int = 16, + min_res: int = 16, + max_res: int = 1024, + hash_table_size: int = 2**19, + features_per_level: int = 2, + hash_init_scale: float = 0.001, + ) -> None: + """_summary_ Args: - x (TensorType[..., "input_dim"]): Input tensor to be encoded + num_levels (int, optional): Number of feature grids. Defaults to 16. + min_res (int, optional): Resolution of smallest feature grid. Defaults to 16. + max_res (int, optional): Resolution of largest feature grid. Defaults to 1024. + hash_table_size (int, optional): Size of hash table. Defaults to 2**19. + features_per_level (int, optional): Number of features per level. Defaults to 2. + hash_init_scale (float, optional): Value to initialize hash grid. Defaults to 0.001. + """ + super().__init__(in_dim=3) + self.num_levels = num_levels + self.features_per_level = features_per_level + self.hash_table_size = hash_table_size + + levels = torch.arange(num_levels) + growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1)) + self.scalings = torch.floor(min_res * growth_factor**levels) + + self.hash_offset = levels * hash_table_size + self.hash_table = torch.rand(size=(hash_table_size * num_levels, features_per_level)) * 2 - 1 + self.hash_table *= hash_init_scale + self.hash_table = torch.nn.Parameter(self.hash_table) + + def get_out_dim(self) -> int: + return self.num_levels * self.features_per_level + + def hash_fn(self, in_tensor: TensorType[..., "num_levels", 3]) -> TensorType[..., "num_levels"]: + """Hash tensor using method described in Instant-NGP + + Args: + in_tensor (TensorType[..., "num_levels", 3]): Tensor to be hashed Returns: - TensorType[..., "output_dim"]: A encoded input tensor + TensorType[..., "num_levels"]: Hashed tensor """ - raise NotImplementedError + in_tensor = in_tensor * torch.tensor([1, 2654435761, 805459861]).to(in_tensor.device) + x = torch.bitwise_xor(in_tensor[..., 0], in_tensor[..., 1]) + x = torch.bitwise_xor(x, in_tensor[..., 2]) + x %= self.hash_table_size + x += self.hash_offset.to(x.device) + return x + + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + assert in_tensor.shape[-1] == 3 + + in_tensor = in_tensor[..., None, :] # [..., 1, 3] + scaled = in_tensor * self.scalings.view(1, 1, -1, 1).to(in_tensor.device) # [..., L, 3] + scaled_c = torch.ceil(scaled).type(torch.int32) + scaled_f = torch.floor(scaled).type(torch.int32) + + offset = scaled - scaled_f + + hashed_0 = self.hash_fn(scaled_c) # [..., num_levels] + hashed_1 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1)) + hashed_2 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_f[..., 1:2], scaled_c[..., 2:3]], dim=-1)) + hashed_3 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_c[..., 2:3]], dim=-1)) + hashed_4 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1)) + hashed_5 = self.hash_fn(torch.cat([scaled_c[..., 0:1], scaled_f[..., 1:2], scaled_f[..., 2:3]], dim=-1)) + hashed_6 = self.hash_fn(scaled_f) + hashed_7 = self.hash_fn(torch.cat([scaled_f[..., 0:1], scaled_c[..., 1:2], scaled_f[..., 2:3]], dim=-1)) + + f_0 = self.hash_table[hashed_0] # [..., num_levels, features_per_level] + f_1 = self.hash_table[hashed_1] + f_2 = self.hash_table[hashed_2] + f_3 = self.hash_table[hashed_3] + f_4 = self.hash_table[hashed_4] + f_5 = self.hash_table[hashed_5] + f_6 = self.hash_table[hashed_6] + f_7 = self.hash_table[hashed_7] + + f_03 = f_0 * offset[..., 0:1] + f_3 * (1 - offset[..., 0:1]) + f_12 = f_1 * offset[..., 0:1] + f_2 * (1 - offset[..., 0:1]) + f_56 = f_5 * offset[..., 0:1] + f_6 * (1 - offset[..., 0:1]) + f_47 = f_4 * offset[..., 0:1] + f_7 * (1 - offset[..., 0:1]) + + f0312 = f_03 * offset[..., 1:2] + f_12 * (1 - offset[..., 1:2]) + f4756 = f_47 * offset[..., 1:2] + f_56 * (1 - offset[..., 1:2]) + + encoded_value = f0312 * offset[..., 2:3] + f4756 * ( + 1 - offset[..., 2:3] + ) # [..., num_levels, features_per_level] + + return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level] diff --git a/mattport/nerf/graph/instant_ngp.py b/mattport/nerf/graph/instant_ngp.py index 5461d0b40..09d0836db 100644 --- a/mattport/nerf/graph/instant_ngp.py +++ b/mattport/nerf/graph/instant_ngp.py @@ -1,4 +1,161 @@ """ -Instant NGP implementation. -TODO: +Implementation of Instant NGP. """ + + +from typing import Dict, List + +import torch +from torch import nn +from torch.nn import Parameter +from torchtyping import TensorType + +from mattport.nerf.field_modules.encoding import HashEncoding, NeRFEncoding +from mattport.nerf.field_modules.field_heads import DensityFieldHead, FieldHeadNames, RGBFieldHead +from mattport.nerf.field_modules.mlp import MLP +from mattport.nerf.field_modules.ray_generator import RayGenerator +from mattport.nerf.graph.base import Graph +from mattport.nerf.loss import MSELoss +from mattport.nerf.renderers import RGBRenderer +from mattport.nerf.sampler import PDFSampler, UniformSampler # pylint: disable=unused-import +from mattport.structures.rays import RaySamples + + +class NGPField(nn.Module): + """NeRF module""" + + def __init__(self, num_layers=3, layer_width=64) -> None: + super().__init__() + self.num_layers = num_layers + self.layer_width = layer_width + + self.encoding_xyz = HashEncoding() + self.encoding_dir = NeRFEncoding( + in_dim=3, num_frequencies=4, min_freq_exp=0.0, max_freq_exp=4.0, include_input=True + ) + self.mlp_base = MLP( + in_dim=self.encoding_xyz.get_out_dim(), + out_dim=self.layer_width, + num_layers=self.num_layers, + layer_width=self.layer_width, + activation=nn.ReLU(), + ) + self.mlp_rgb = MLP( + in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(), + out_dim=self.layer_width // 2, + num_layers=1, + activation=nn.ReLU(), + ) + self.field_output_rgb = RGBFieldHead(in_dim=self.mlp_rgb.get_out_dim()) + self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim()) + + def forward(self, ray_samples: RaySamples): + """Evaluates the field at points along the ray + Args: + xyz: () + # TODO(ethan): change the input to be something more abstracted + e.g., a FieldInput structure + """ + positions = ray_samples.positions + directions = ray_samples.directions + encoded_xyz = self.encoding_xyz(positions) + encoded_dir = self.encoding_dir(directions) + base_mlp_out = self.mlp_base(encoded_xyz) + rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1)) + + field_rgb_output = self.field_output_rgb(rgb_mlp_out) + field_density_out = self.field_output_density(base_mlp_out) + + field_outputs = {} + field_outputs.update(field_rgb_output) + field_outputs.update(field_density_out) + return field_outputs + + +class NGPGraph(Graph): + """Vanilla NeRF graph""" + + def __init__( + self, + intrinsics=None, + camera_to_world=None, + near_plane=1.0, + far_plane=6.0, + num_coarse_samples=64, + num_importance_samples=128, + **kwargs + ) -> None: + self.near_plane = near_plane + self.far_plane = far_plane + self.num_coarse_samples = num_coarse_samples + self.num_importance_samples = num_importance_samples + super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world, **kwargs) + + def populate_modules(self): + # ray generator + self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world) + + # samplers + self.sampler_uniform = UniformSampler( + near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_coarse_samples + ) + self.sampler_pdf = PDFSampler(num_samples=self.num_importance_samples) + + # field + self.field_coarse = NGPField() + self.field_fine = NGPField() + + # renderers + self.renderer_rgb = RGBRenderer() + + # losses + self.rgb_loss = MSELoss() + + def get_param_groups(self) -> Dict[str, List[Parameter]]: + """Obtain the parameter groups for the optimizers + + Returns: + Dict[str, List[Parameter]]: Mapping of different parameter groups + """ + param_groups = {} + param_groups["fields"] = list(self.field_coarse.parameters()) + list(self.field_fine.parameters()) + return param_groups + + def forward(self, ray_indices: TensorType["num_rays", 3]): + """Takes in the ray indices and renders out values.""" + # get the rays: + ray_bundle = self.ray_generator.forward(ray_indices) # RayBundle + # coarse network: + uniform_ray_samples = self.sampler_uniform(ray_bundle) # RaySamples + + # time coarse network # xyz -> asd + coarse_field_outputs = self.field_coarse(uniform_ray_samples) # FieldOutputs + # time end coarse network + + coarse_weights = uniform_ray_samples.get_weights(coarse_field_outputs[FieldHeadNames.DENSITY]) + + coarse_renderer_outputs = self.renderer_rgb( + rgb=coarse_field_outputs[FieldHeadNames.RGB], + weights=coarse_weights, + ) # RendererOutputs + # fine network: + pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples + fine_field_outputs = self.field_fine(pdf_ray_samples) # FieldOutputs + + fine_weights = pdf_ray_samples.get_weights(fine_field_outputs[FieldHeadNames.DENSITY]) + + fine_renderer_outputs = self.renderer_rgb( + rgb=fine_field_outputs[FieldHeadNames.RGB], + weights=fine_weights, + ) # RendererOutputs + # outputs: + outputs = {"rgb_coarse": coarse_renderer_outputs.rgb, "rgb_fine": fine_renderer_outputs.rgb} + return outputs + + def get_losses(self, batch, graph_outputs): + # batch.pixels # (num_rays, 3) + losses = {} + rgb_loss_coarse = self.rgb_loss(batch["pixels"], graph_outputs["rgb_coarse"]) + rgb_loss_fine = self.rgb_loss(batch["pixels"], graph_outputs["rgb_fine"]) + losses = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} + return losses From 00d9046c8eaf1db0f541c6ed7656a3f92b85c292 Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Mon, 2 May 2022 15:10:49 -0700 Subject: [PATCH 2/2] Update lint checks --- .github/workflows/code_checks.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml index b850c9a51..3cfc108c1 100644 --- a/.github/workflows/code_checks.yml +++ b/.github/workflows/code_checks.yml @@ -37,7 +37,8 @@ jobs: --disable=too-many-arguments \ --disable=too-many-locals \ --disable=fixme \ - --disable=too-many-instance-attributes + --disable=too-many-instance-attributes \ + --disable=duplicate-code - name: Test with pytest run: | pytest tests --typeguard-packages=mattport --torchtyping-patch-typeguard --disable-warnings