From 78725bcb34382841e2685644abbe6702d72e399e Mon Sep 17 00:00:00 2001 From: terrance Date: Sat, 3 Jun 2023 04:41:17 +0000 Subject: [PATCH] update fields to use tcnn --- nerfstudio/configs/method_configs.py | 8 +- nerfstudio/fields/dreamfusion_field.py | 143 +++++++++++++++---------- nerfstudio/models/dreamfusion.py | 7 +- 3 files changed, 96 insertions(+), 62 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 11a1e8b88..14ab619a2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -481,15 +481,15 @@ method_configs["dreamfusion"] = TrainerConfig( ), optimizers={ "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" + warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear" ), }, "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - warmup_steps=2000, lr_final=1e-6, max_steps=20000, ramp="linear" + warmup_steps=4000, lr_final=5e-7, max_steps=20000, ramp="linear" ), }, }, diff --git a/nerfstudio/fields/dreamfusion_field.py b/nerfstudio/fields/dreamfusion_field.py index a42bd4697..bc7cb9136 100644 --- a/nerfstudio/fields/dreamfusion_field.py +++ b/nerfstudio/fields/dreamfusion_field.py @@ -17,22 +17,24 @@ Field for dreamfusion model """ -from typing import Optional, Tuple, Dict +from typing import Optional, Tuple, Dict, Literal import numpy as np import torch from torch.nn.parameter import Parameter from jaxtyping import Float -from torch import Tensor +from torch import Tensor, nn +from nerfstudio.field_components.mlp import MLP from nerfstudio.cameras.rays import RayBundle, RaySamples from nerfstudio.data.scene_box import SceneBox from nerfstudio.field_components.activations import trunc_exp from nerfstudio.field_components.field_heads import FieldHeadNames +from nerfstudio.field_components.encodings import HashEncoding, SHEncoding from nerfstudio.fields.base_field import Field try: - import tinycudann as tcnn + pass except ImportError: # tinycudann module doesn't exist pass @@ -73,6 +75,7 @@ class DreamFusionField(Field): log2_hashmap_size: int = 19, num_layers_color: int = 3, hidden_dim_color: int = 64, + implementation: Literal["tcnn", "torch"] = "tcnn", ) -> None: super().__init__() @@ -81,64 +84,94 @@ class DreamFusionField(Field): base_res = 16 features_per_level = 2 - growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) + np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1)) - self.position_encoding = tcnn.Encoding( - n_input_dims=3, - encoding_config={"otype": "Frequency", "n_frequencies": 2}, + self.direction_encoding = SHEncoding( + levels=4, + implementation=implementation, ) - self.direction_encoding = tcnn.Encoding( - n_input_dims=3, - encoding_config={ - "otype": "SphericalHarmonics", - "degree": 4, - }, + encoder = HashEncoding( + num_levels=num_levels, + min_res=base_res, + max_res=max_res, + log2_hashmap_size=log2_hashmap_size, + features_per_level=features_per_level, + implementation=implementation, + ) + network = MLP( + in_dim=encoder.get_out_dim(), + num_layers=num_layers, + layer_width=hidden_dim, + out_dim=1 + self.geo_feat_dim, + activation=nn.ReLU(), + out_activation=None, + implementation=implementation, + ) + self.mlp_base = torch.nn.Sequential(encoder, network) + + # self.mlp_base = tcnn.NetworkWithInputEncoding( + # n_input_dims=3, + # n_output_dims=1 + self.geo_feat_dim, + # encoding_config={ + # "otype": "HashGrid", + # "n_levels": num_levels, + # "n_features_per_level": features_per_level, + # "log2_hashmap_size": log2_hashmap_size, + # "base_resolution": base_res, + # "per_level_scale": growth_factor, + # }, + # network_config={ + # "otype": "FullyFusedMLP", + # "activation": "ReLU", + # "output_activation": "None", + # "n_neurons": hidden_dim, + # "n_hidden_layers": num_layers - 1, + # }, + # ) + + # self.mlp_background_color = tcnn.Network( + # n_input_dims=self.direction_encoding.n_output_dims, + # n_output_dims=3, + # network_config={ + # "otype": "FullyFusedMLP", + # "activation": "ReLU", + # "output_activation": "Sigmoid", + # "n_neurons": 32, + # "n_hidden_layers": 1, + # }, + # ) + + self.mlp_background_color = MLP( + in_dim=self.direction_encoding.get_out_dim(), + num_layers=2, + layer_width=32, + out_dim=3, + activation=nn.ReLU(), + out_activation=nn.Sigmoid(), + implementation=implementation, ) - self.mlp_base = tcnn.NetworkWithInputEncoding( - n_input_dims=3, - n_output_dims=1 + self.geo_feat_dim, - encoding_config={ - "otype": "HashGrid", - "n_levels": num_levels, - "n_features_per_level": features_per_level, - "log2_hashmap_size": log2_hashmap_size, - "base_resolution": base_res, - "per_level_scale": growth_factor, - }, - network_config={ - "otype": "FullyFusedMLP", - "activation": "ReLU", - "output_activation": "None", - "n_neurons": hidden_dim, - "n_hidden_layers": num_layers - 1, - }, - ) - - self.mlp_background_color = tcnn.Network( - n_input_dims=self.direction_encoding.n_output_dims, - n_output_dims=3, - network_config={ - "otype": "FullyFusedMLP", - "activation": "ReLU", - "output_activation": "Sigmoid", - "n_neurons": 32, - "n_hidden_layers": 1, - }, - ) - - self.mlp_head = tcnn.Network( - n_input_dims=self.geo_feat_dim, - n_output_dims=3, - network_config={ - "otype": "FullyFusedMLP", - "activation": "ReLU", - "output_activation": "Sigmoid", - "n_neurons": hidden_dim_color, - "n_hidden_layers": num_layers_color - 1, - }, + self.mlp_head = MLP( + in_dim=self.geo_feat_dim, + num_layers=num_layers_color, + layer_width=hidden_dim_color, + out_dim=3, + activation=nn.ReLU(), + out_activation=nn.Sigmoid(), + implementation=implementation, ) + # self.mlp_head = tcnn.Network( + # n_input_dims=self.geo_feat_dim, + # n_output_dims=3, + # network_config={ + # "otype": "FullyFusedMLP", + # "activation": "ReLU", + # "output_activation": "Sigmoid", + # "n_neurons": hidden_dim_color, + # "n_hidden_layers": num_layers_color - 1, + # }, + # ) def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]: """Computes and returns the densities.""" diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 3bd163416..d6bee2906 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -66,7 +66,7 @@ class DreamFusionModelConfig(ModelConfig): _target: Type = field(default_factory=lambda: DreamFusionModel) """target class to instantiate""" - prompt: str = "a high quality photo of a tree frog" + prompt: str = "a high quality zoomed out photo of an orange tabby cat sitting on a box" """prompt for stable dreamfusion""" orientation_loss_mult: Tuple[float, float] = (0.01, 100.0) @@ -159,6 +159,8 @@ class DreamFusionModelConfig(ModelConfig): """diffusion model for SDS loss""" sd_version: str = "1-5" """model version when using stable diffusion""" + implementation: Literal["tcnn", "torch"] = "tcnn" + """Which implementation to use for the model.""" class DreamFusionModel(Model): @@ -229,8 +231,7 @@ class DreamFusionModel(Model): for i in range(num_prop_nets): prop_net_args = self.config.proposal_net_args_list[min(i, len(self.config.proposal_net_args_list) - 1)] network = HashMLPDensityField( - self.scene_box.aabb, - **prop_net_args, + self.scene_box.aabb, **prop_net_args, implementation=self.config.implementation ) self.proposal_networks.append(network) self.density_fns.extend([network.density_fn for network in self.proposal_networks])