From 8fad92c39af4e2d992dfdf950d4cff4f4478c3c9 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Thu, 15 Jun 2023 11:05:53 -0700 Subject: [PATCH] nerfacto-huge (#2003) An even larger version of nerfacto with scaled hashgrid resolution, MLP sizes, batch sizes, appearance embeddings, proposal sampler resolution, and more. Also includes some minor tweaks to nerfacto-big configs --- nerfstudio/configs/method_configs.py | 70 +++++++++++++++++++++------ nerfstudio/engine/trainer.py | 6 +-- nerfstudio/fields/nerfacto_field.py | 14 +++--- nerfstudio/model_components/losses.py | 3 +- nerfstudio/models/nerfacto.py | 25 ++++------ 5 files changed, 77 insertions(+), 41 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 4d0c3dad1..6c116e645 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -26,19 +26,12 @@ import tyro from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.configs.base_config import ViewerConfig from nerfstudio.configs.external_methods import get_external_methods -from nerfstudio.data.datamanagers.base_datamanager import ( - VanillaDataManager, - VanillaDataManagerConfig, -) +from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager, VanillaDataManagerConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig -from nerfstudio.data.dataparsers.instant_ngp_dataparser import ( - InstantNGPDataParserConfig, -) +from nerfstudio.data.dataparsers.instant_ngp_dataparser import InstantNGPDataParserConfig from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig -from nerfstudio.data.dataparsers.phototourism_dataparser import ( - PhototourismDataParserConfig, -) +from nerfstudio.data.dataparsers.phototourism_dataparser import PhototourismDataParserConfig from nerfstudio.data.dataparsers.sdfstudio_dataparser import SDFStudioDataParserConfig from nerfstudio.data.dataparsers.sitcoms3d_dataparser import Sitcoms3DDataParserConfig from nerfstudio.data.datasets.depth_dataset import DepthDataset @@ -124,7 +117,7 @@ method_configs["nerfacto-big"] = TrainerConfig( pipeline=VanillaPipelineConfig( datamanager=VanillaDataManagerConfig( dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, + train_num_rays_per_batch=8192, eval_num_rays_per_batch=4096, camera_optimizer=CameraOptimizerConfig( mode="SO3xR3", @@ -137,8 +130,9 @@ method_configs["nerfacto-big"] = TrainerConfig( num_proposal_samples_per_ray=(512, 256), hidden_dim=128, hidden_dim_color=128, - hidden_dim_transient=128, - max_res=3000, + appearance_embed_dim=128, + base_res=32, + max_res=4096, proposal_weights_anneal_max_num_iters=5000, log2_hashmap_size=21, ), @@ -150,7 +144,55 @@ method_configs["nerfacto-big"] = TrainerConfig( }, "fields": { "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=100000), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", +) +method_configs["nerfacto-huge"] = TrainerConfig( + method_name="nerfacto", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=100000, + mixed_precision=True, + pipeline=VanillaPipelineConfig( + datamanager=VanillaDataManagerConfig( + dataparser=NerfstudioDataParserConfig(downscale_factor=1), + train_num_rays_per_batch=16384, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", + optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), + scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), + ), + ), + model=NerfactoModelConfig( + eval_num_rays_per_chunk=1 << 15, + num_nerf_samples_per_ray=64, + num_proposal_samples_per_ray=(512, 512), + proposal_net_args_list=[ + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False}, + ], + hidden_dim=256, + hidden_dim_color=256, + appearance_embed_dim=256, + features_per_level=4, + base_res=32, + max_res=8192, + proposal_weights_anneal_max_num_iters=5000, + log2_hashmap_size=21, + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 713263234..afc0d94f3 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -453,10 +453,8 @@ class Trainer: """ self.optimizers.zero_grad_all() - - device_type: str = self.device.split(":")[0] if "cuda" in self.device else "cpu" - - with torch.autocast(device_type=device_type, enabled=self.mixed_precision): + cpu_or_cuda_str: str = self.device.split(":")[0] + with torch.autocast(device_type=cpu_or_cuda_str, enabled=self.mixed_precision): _, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step=step) loss = functools.reduce(torch.add, loss_dict.values()) self.grad_scaler.scale(loss).backward() # type: ignore diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py index 898d344d3..1c35b7121 100644 --- a/nerfstudio/fields/nerfacto_field.py +++ b/nerfstudio/fields/nerfacto_field.py @@ -76,10 +76,12 @@ class NerfactoField(Field): hidden_dim: int = 64, geo_feat_dim: int = 15, num_levels: int = 16, + base_res: int = 16, max_res: int = 2048, log2_hashmap_size: int = 19, num_layers_color: int = 3, num_layers_transient: int = 2, + features_per_level: int = 2, hidden_dim_color: int = 64, hidden_dim_transient: int = 64, appearance_embedding_dim: int = 32, @@ -111,9 +113,7 @@ class NerfactoField(Field): self.use_semantics = use_semantics self.use_pred_normals = use_pred_normals self.pass_semantic_gradients = pass_semantic_gradients - - base_res: int = 16 - features_per_level: int = 2 + self.base_res = base_res self.direction_encoding = SHEncoding( levels=4, @@ -124,7 +124,7 @@ class NerfactoField(Field): in_dim=3, num_frequencies=2, min_freq_exp=0, max_freq_exp=2 - 1, implementation=implementation ) - encoder = HashEncoding( + self.mlp_base_grid = HashEncoding( num_levels=num_levels, min_res=base_res, max_res=max_res, @@ -132,8 +132,8 @@ class NerfactoField(Field): features_per_level=features_per_level, implementation=implementation, ) - network = MLP( - in_dim=encoder.get_out_dim(), + self.mlp_base_mlp = MLP( + in_dim=self.mlp_base_grid.get_out_dim(), num_layers=num_layers, layer_width=hidden_dim, out_dim=1 + self.geo_feat_dim, @@ -141,7 +141,7 @@ class NerfactoField(Field): out_activation=None, implementation=implementation, ) - self.mlp_base = torch.nn.Sequential(encoder, network) + self.mlp_base = torch.nn.Sequential(self.mlp_base_grid, self.mlp_base_mlp) # transients if self.use_transient_embedding: diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 604f306e6..44fefafd1 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -542,7 +542,8 @@ class _GradientScaler(torch.autograd.Function): # typing: ignore def scale_gradients_by_distance_squared( - field_outputs: Dict[FieldHeadNames, torch.Tensor], ray_samples: RaySamples + field_outputs: Dict[FieldHeadNames, torch.Tensor], + ray_samples: RaySamples, ) -> Dict[FieldHeadNames, torch.Tensor]: """ Scale gradients by the ray distance to the pixel diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 80ead5951..75e798873 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -29,11 +29,7 @@ from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle, RaySamples -from nerfstudio.engine.callbacks import ( - TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation, -) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.field_components.field_heads import FieldHeadNames from nerfstudio.field_components.spatial_distortions import SceneContraction from nerfstudio.fields.density_fields import HashMLPDensityField @@ -46,16 +42,8 @@ from nerfstudio.model_components.losses import ( pred_normal_loss, scale_gradients_by_distance_squared, ) -from nerfstudio.model_components.ray_samplers import ( - ProposalNetworkSampler, - UniformSampler, -) -from nerfstudio.model_components.renderers import ( - AccumulationRenderer, - DepthRenderer, - NormalsRenderer, - RGBRenderer, -) +from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler, UniformSampler +from nerfstudio.model_components.renderers import AccumulationRenderer, DepthRenderer, NormalsRenderer, RGBRenderer from nerfstudio.model_components.scene_colliders import NearFarCollider from nerfstudio.model_components.shaders import NormalsShader from nerfstudio.models.base_model import Model, ModelConfig @@ -81,10 +69,14 @@ class NerfactoModelConfig(ModelConfig): """Dimension of hidden layers for transient network""" num_levels: int = 16 """Number of levels of the hashmap for the base mlp.""" + base_res: int = 16 + """Resolution of the base grid for the hasgrid.""" max_res: int = 2048 """Maximum resolution of the hashmap for the base mlp.""" log2_hashmap_size: int = 19 """Size of the hashmap for the base mlp""" + features_per_level: int = 2 + """How many hashgrid features per level""" num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96) """Number of samples per ray for each proposal network.""" num_nerf_samples_per_ray: int = 48 @@ -132,6 +124,8 @@ class NerfactoModelConfig(ModelConfig): """Use gradient scaler where the gradients are lower for points closer to the camera.""" implementation: Literal["tcnn", "torch"] = "tcnn" """Which implementation to use for the model.""" + appearance_embed_dim: int = 32 + """Dimension of the appearance embedding.""" class NerfactoModel(Model): @@ -165,6 +159,7 @@ class NerfactoModel(Model): num_images=self.num_train_data, use_pred_normals=self.config.predict_normals, use_average_appearance_embedding=self.config.use_average_appearance_embedding, + appearance_embedding_dim=self.config.appearance_embed_dim, implementation=self.config.implementation, )