mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 19:33:46 +08:00
rip out camera optimization from datamanager and move to nerfacto.py
This commit is contained in:
Justin committed
1 parent
1182b18b48
commit
ae5ef15bed
7 files changed
+497
-548
No files matched your search
@@ -30,11 +30,7 @@ from typing_extensions import assert_never
|
||||
|
||||
from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3
|
||||
from nerfstudio.configs.base_config import InstantiateConfig
|
||||
from nerfstudio.engine.optimizers import AdamOptimizerConfig, OptimizerConfig
|
||||
from nerfstudio.engine.schedulers import (
|
||||
ExponentialDecaySchedulerConfig,
|
||||
SchedulerConfig,
|
||||
)
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from nerfstudio.utils import poses as pose_utils
|
||||
|
||||
|
||||
@@ -47,22 +43,6 @@ class CameraOptimizerConfig(InstantiateConfig):
|
||||
mode: Literal["off", "SO3xR3", "SE3"] = "off"
|
||||
"""Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
|
||||
|
||||
position_noise_std: float = 0.0
|
||||
"""Noise to add to initial positions. Useful for debugging."""
|
||||
|
||||
orientation_noise_std: float = 0.0
|
||||
"""Noise to add to initial orientations. Useful for debugging."""
|
||||
|
||||
optimizer: OptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15)
|
||||
"""ADAM parameters for camera optimization."""
|
||||
|
||||
scheduler: SchedulerConfig = ExponentialDecaySchedulerConfig(max_steps=10000)
|
||||
"""Learning rate scheduler for camera optimizer.."""
|
||||
|
||||
param_group: tyro.conf.Suppress[str] = "camera_opt"
|
||||
"""Name of the parameter group used for pose optimization. Can be any string that doesn't conflict with other
|
||||
groups."""
|
||||
|
||||
|
||||
class CameraOptimizer(nn.Module):
|
||||
"""Layer that modifies camera poses to be optimized as well as the field during training."""
|
||||
@@ -89,16 +69,6 @@ class CameraOptimizer(nn.Module):
|
||||
else:
|
||||
assert_never(self.config.mode)
|
||||
|
||||
# Initialize pose noise; useful for debugging.
|
||||
if config.position_noise_std != 0.0 or config.orientation_noise_std != 0.0:
|
||||
assert config.position_noise_std >= 0.0 and config.orientation_noise_std >= 0.0
|
||||
std_vector = torch.tensor(
|
||||
[config.position_noise_std] * 3 + [config.orientation_noise_std] * 3, device=device
|
||||
)
|
||||
self.pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=device), std_vector))
|
||||
else:
|
||||
self.pose_noise = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
indices: Int[Tensor, "num_cameras"],
|
||||
@@ -122,12 +92,16 @@ class CameraOptimizer(nn.Module):
|
||||
else:
|
||||
assert_never(self.config.mode)
|
||||
|
||||
# Apply initial pose noise.
|
||||
if self.pose_noise is not None:
|
||||
outputs.append(self.pose_noise[indices, :, :])
|
||||
|
||||
# Return: identity if no transforms are needed, otherwise multiply transforms together.
|
||||
if len(outputs) == 0:
|
||||
# Note that using repeat() instead of tile() here would result in unnecessary copies.
|
||||
return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1)
|
||||
return functools.reduce(pose_utils.multiply, outputs)
|
||||
|
||||
def apply_to_raybundle(self, raybundle: RayBundle) -> None:
|
||||
"""
|
||||
Apply the pose correction to the raybundle
|
||||
"""
|
||||
correction_matrices = self(raybundle.camera_indices.squeeze())
|
||||
raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
|
||||
raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
|
||||
@@ -87,486 +87,492 @@ method_configs["nerfacto"] = TrainerConfig(
|
||||
dataparser=NerfstudioDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3",
|
||||
optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2),
|
||||
scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000),
|
||||
),
|
||||
),
|
||||
model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
method_configs["nerfacto-big"] = 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(),
|
||||
train_num_rays_per_batch=8192,
|
||||
eval_num_rays_per_batch=4096,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3",
|
||||
optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3),
|
||||
),
|
||||
),
|
||||
model=NerfactoModelConfig(
|
||||
eval_num_rays_per_chunk=1 << 15,
|
||||
num_nerf_samples_per_ray=128,
|
||||
num_proposal_samples_per_ray=(512, 256),
|
||||
hidden_dim=128,
|
||||
hidden_dim_color=128,
|
||||
appearance_embed_dim=128,
|
||||
base_res=32,
|
||||
max_res=4096,
|
||||
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),
|
||||
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(),
|
||||
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=32,
|
||||
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),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["depth-nerfacto"] = TrainerConfig(
|
||||
method_name="depth-nerfacto",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
_target=VanillaDataManager[DepthDataset],
|
||||
dataparser=NerfstudioDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
),
|
||||
),
|
||||
model=DepthNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["volinga"] = TrainerConfig(
|
||||
method_name="volinga",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
dataparser=NerfstudioDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
),
|
||||
),
|
||||
model=NerfactoModelConfig(
|
||||
eval_num_rays_per_chunk=1 << 15,
|
||||
hidden_dim=32,
|
||||
hidden_dim_color=32,
|
||||
hidden_dim_transient=32,
|
||||
num_nerf_samples_per_ray=24,
|
||||
proposal_net_args_list=[
|
||||
{"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True},
|
||||
{"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True},
|
||||
],
|
||||
camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"),
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["instant-ngp"] = TrainerConfig(
|
||||
method_name="instant-ngp",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=DynamicBatchPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
dataparser=NerfstudioDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
),
|
||||
model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
}
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 12),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
|
||||
method_configs["instant-ngp-bounded"] = TrainerConfig(
|
||||
method_name="instant-ngp-bounded",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=DynamicBatchPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192),
|
||||
model=InstantNGPModelConfig(
|
||||
eval_num_rays_per_chunk=8192,
|
||||
grid_levels=1,
|
||||
alpha_thre=0.0,
|
||||
cone_angle=0.0,
|
||||
disable_scene_contraction=True,
|
||||
near_plane=0.01,
|
||||
background_color="black",
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
}
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 12),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
|
||||
method_configs["mipnerf"] = TrainerConfig(
|
||||
method_name="mipnerf",
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024),
|
||||
model=VanillaModelConfig(
|
||||
_target=MipNerfModel,
|
||||
loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0},
|
||||
num_coarse_samples=128,
|
||||
num_importance_samples=128,
|
||||
eval_num_rays_per_chunk=1024,
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
"scheduler": None,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
method_configs["semantic-nerfw"] = TrainerConfig(
|
||||
method_name="semantic-nerfw",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
_target=VanillaDataManager[SemanticDataset],
|
||||
dataparser=Sitcoms3DDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=8192,
|
||||
),
|
||||
model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 16),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["vanilla-nerf"] = TrainerConfig(
|
||||
method_name="vanilla-nerf",
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
dataparser=BlenderDataParserConfig(),
|
||||
),
|
||||
model=VanillaModelConfig(_target=NeRFModel),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
"scheduler": None,
|
||||
},
|
||||
"temporal_distortion": {
|
||||
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
method_configs["tensorf"] = TrainerConfig(
|
||||
method_name="tensorf",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=False,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
dataparser=BlenderDataParserConfig(),
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
),
|
||||
model=TensoRFModelConfig(
|
||||
regularization="tv",
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=0.001),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
|
||||
},
|
||||
"encodings": {
|
||||
"optimizer": AdamOptimizerConfig(lr=0.02),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000),
|
||||
"camera_opt": {
|
||||
"optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2),
|
||||
"scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000),
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["dnerf"] = TrainerConfig(
|
||||
method_name="dnerf",
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()),
|
||||
model=VanillaModelConfig(
|
||||
_target=NeRFModel,
|
||||
enable_temporal_distortion=True,
|
||||
temporal_distortion_params={"kind": TemporalDistortionKind.DNERF},
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
"scheduler": None,
|
||||
},
|
||||
"temporal_distortion": {
|
||||
"optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
)
|
||||
# **For now I commented everything out since the datamanager config breaks **
|
||||
|
||||
method_configs["phototourism"] = TrainerConfig(
|
||||
method_name="phototourism",
|
||||
steps_per_eval_batch=500,
|
||||
steps_per_save=2000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=True,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto
|
||||
train_num_rays_per_batch=4096,
|
||||
eval_num_rays_per_batch=4096,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
),
|
||||
# Large dataset, so using prior values from VariableResDataManager.
|
||||
train_num_images_to_sample_from=40,
|
||||
train_num_times_to_repeat_images=100,
|
||||
eval_num_images_to_sample_from=40,
|
||||
eval_num_times_to_repeat_images=100,
|
||||
),
|
||||
model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": None,
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
|
||||
method_configs["neus"] = TrainerConfig(
|
||||
method_name="neus",
|
||||
steps_per_eval_image=500,
|
||||
steps_per_eval_batch=5000,
|
||||
steps_per_save=20000,
|
||||
steps_per_eval_all_images=1000000, # set to a very large number so we don't eval with all images
|
||||
max_num_iterations=100000,
|
||||
mixed_precision=False,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
_target=VanillaDataManager[SDFDataset],
|
||||
dataparser=SDFStudioDataParserConfig(),
|
||||
train_num_rays_per_batch=1024,
|
||||
eval_num_rays_per_batch=1024,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
),
|
||||
),
|
||||
model=NeuSModelConfig(eval_num_rays_per_chunk=1024),
|
||||
),
|
||||
optimizers={
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
"scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
|
||||
},
|
||||
"field_background": {
|
||||
"optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
"scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
# method_configs["nerfacto-big"] = 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(),
|
||||
# train_num_rays_per_batch=8192,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="SO3xR3",
|
||||
# optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3),
|
||||
# ),
|
||||
# ),
|
||||
# model=NerfactoModelConfig(
|
||||
# eval_num_rays_per_chunk=1 << 15,
|
||||
# num_nerf_samples_per_ray=128,
|
||||
# num_proposal_samples_per_ray=(512, 256),
|
||||
# hidden_dim=128,
|
||||
# hidden_dim_color=128,
|
||||
# appearance_embed_dim=128,
|
||||
# base_res=32,
|
||||
# max_res=4096,
|
||||
# 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),
|
||||
# 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(),
|
||||
# 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=32,
|
||||
# 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),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
method_configs["neus-facto"] = TrainerConfig(
|
||||
method_name="neus-facto",
|
||||
steps_per_eval_image=5000,
|
||||
steps_per_eval_batch=5000,
|
||||
steps_per_save=2000,
|
||||
steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images
|
||||
max_num_iterations=20001,
|
||||
mixed_precision=False,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=VanillaDataManagerConfig(
|
||||
_target=VanillaDataManager[SDFDataset],
|
||||
dataparser=SDFStudioDataParserConfig(),
|
||||
train_num_rays_per_batch=2048,
|
||||
eval_num_rays_per_batch=2048,
|
||||
camera_optimizer=CameraOptimizerConfig(
|
||||
mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
),
|
||||
),
|
||||
model=NeuSFactoModelConfig(
|
||||
# proposal network allows for significantly smaller sdf/color network
|
||||
sdf_field=SDFFieldConfig(
|
||||
use_grid_feature=True,
|
||||
num_layers=2,
|
||||
num_layers_color=2,
|
||||
hidden_dim=256,
|
||||
bias=0.5,
|
||||
beta_init=0.8,
|
||||
use_appearance_embedding=False,
|
||||
),
|
||||
background_model="none",
|
||||
eval_num_rays_per_chunk=2048,
|
||||
),
|
||||
),
|
||||
optimizers={
|
||||
"proposal_networks": {
|
||||
"optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
"scheduler": MultiStepSchedulerConfig(max_steps=20001, milestones=(10000, 1500, 18000)),
|
||||
},
|
||||
"fields": {
|
||||
"optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
"scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),
|
||||
},
|
||||
"field_background": {
|
||||
"optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
"scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),
|
||||
},
|
||||
},
|
||||
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
vis="viewer",
|
||||
)
|
||||
# method_configs["depth-nerfacto"] = TrainerConfig(
|
||||
# method_name="depth-nerfacto",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# _target=VanillaDataManager[DepthDataset],
|
||||
# dataparser=NerfstudioDataParserConfig(),
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
# ),
|
||||
# ),
|
||||
# model=DepthNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "proposal_networks": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["volinga"] = TrainerConfig(
|
||||
# method_name="volinga",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# dataparser=NerfstudioDataParserConfig(),
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
# ),
|
||||
# ),
|
||||
# model=NerfactoModelConfig(
|
||||
# eval_num_rays_per_chunk=1 << 15,
|
||||
# hidden_dim=32,
|
||||
# hidden_dim_color=32,
|
||||
# hidden_dim_transient=32,
|
||||
# num_nerf_samples_per_ray=24,
|
||||
# proposal_net_args_list=[
|
||||
# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True},
|
||||
# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True},
|
||||
# ],
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "proposal_networks": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["instant-ngp"] = TrainerConfig(
|
||||
# method_name="instant-ngp",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=DynamicBatchPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# dataparser=NerfstudioDataParserConfig(),
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# ),
|
||||
# model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
# }
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 12),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
|
||||
# method_configs["instant-ngp-bounded"] = TrainerConfig(
|
||||
# method_name="instant-ngp-bounded",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=DynamicBatchPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192),
|
||||
# model=InstantNGPModelConfig(
|
||||
# eval_num_rays_per_chunk=8192,
|
||||
# grid_levels=1,
|
||||
# alpha_thre=0.0,
|
||||
# cone_angle=0.0,
|
||||
# disable_scene_contraction=True,
|
||||
# near_plane=0.01,
|
||||
# background_color="black",
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
|
||||
# }
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 12),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
|
||||
# method_configs["mipnerf"] = TrainerConfig(
|
||||
# method_name="mipnerf",
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024),
|
||||
# model=VanillaModelConfig(
|
||||
# _target=MipNerfModel,
|
||||
# loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0},
|
||||
# num_coarse_samples=128,
|
||||
# num_importance_samples=128,
|
||||
# eval_num_rays_per_chunk=1024,
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
# "scheduler": None,
|
||||
# }
|
||||
# },
|
||||
# )
|
||||
|
||||
# method_configs["semantic-nerfw"] = TrainerConfig(
|
||||
# method_name="semantic-nerfw",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# _target=VanillaDataManager[SemanticDataset],
|
||||
# dataparser=Sitcoms3DDataParserConfig(),
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=8192,
|
||||
# ),
|
||||
# model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "proposal_networks": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 16),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["vanilla-nerf"] = TrainerConfig(
|
||||
# method_name="vanilla-nerf",
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# dataparser=BlenderDataParserConfig(),
|
||||
# ),
|
||||
# model=VanillaModelConfig(_target=NeRFModel),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "temporal_distortion": {
|
||||
# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# )
|
||||
|
||||
# method_configs["tensorf"] = TrainerConfig(
|
||||
# method_name="tensorf",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=False,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# dataparser=BlenderDataParserConfig(),
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# ),
|
||||
# model=TensoRFModelConfig(
|
||||
# regularization="tv",
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=0.001),
|
||||
# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000),
|
||||
# },
|
||||
# "encodings": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=0.02),
|
||||
# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000),
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["dnerf"] = TrainerConfig(
|
||||
# method_name="dnerf",
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()),
|
||||
# model=VanillaModelConfig(
|
||||
# _target=NeRFModel,
|
||||
# enable_temporal_distortion=True,
|
||||
# temporal_distortion_params={"kind": TemporalDistortionKind.DNERF},
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "temporal_distortion": {
|
||||
# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# )
|
||||
|
||||
# method_configs["phototourism"] = TrainerConfig(
|
||||
# method_name="phototourism",
|
||||
# steps_per_eval_batch=500,
|
||||
# steps_per_save=2000,
|
||||
# max_num_iterations=30000,
|
||||
# mixed_precision=True,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto
|
||||
# train_num_rays_per_batch=4096,
|
||||
# eval_num_rays_per_batch=4096,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
# ),
|
||||
# # Large dataset, so using prior values from VariableResDataManager.
|
||||
# train_num_images_to_sample_from=40,
|
||||
# train_num_times_to_repeat_images=100,
|
||||
# eval_num_images_to_sample_from=40,
|
||||
# eval_num_times_to_repeat_images=100,
|
||||
# ),
|
||||
# model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "proposal_networks": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": None,
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["neus"] = TrainerConfig(
|
||||
# method_name="neus",
|
||||
# steps_per_eval_image=500,
|
||||
# steps_per_eval_batch=5000,
|
||||
# steps_per_save=20000,
|
||||
# steps_per_eval_all_images=1000000, # set to a very large number so we don't eval with all images
|
||||
# max_num_iterations=100000,
|
||||
# mixed_precision=False,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# _target=VanillaDataManager[SDFDataset],
|
||||
# dataparser=SDFStudioDataParserConfig(),
|
||||
# train_num_rays_per_batch=1024,
|
||||
# eval_num_rays_per_batch=1024,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
# ),
|
||||
# ),
|
||||
# model=NeuSModelConfig(eval_num_rays_per_chunk=1024),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
# "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
|
||||
# },
|
||||
# "field_background": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
# "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000),
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
# method_configs["neus-facto"] = TrainerConfig(
|
||||
# method_name="neus-facto",
|
||||
# steps_per_eval_image=5000,
|
||||
# steps_per_eval_batch=5000,
|
||||
# steps_per_save=2000,
|
||||
# steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images
|
||||
# max_num_iterations=20001,
|
||||
# mixed_precision=False,
|
||||
# pipeline=VanillaPipelineConfig(
|
||||
# datamanager=VanillaDataManagerConfig(
|
||||
# _target=VanillaDataManager[SDFDataset],
|
||||
# dataparser=SDFStudioDataParserConfig(),
|
||||
# train_num_rays_per_batch=2048,
|
||||
# eval_num_rays_per_batch=2048,
|
||||
# camera_optimizer=CameraOptimizerConfig(
|
||||
# mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2)
|
||||
# ),
|
||||
# ),
|
||||
# model=NeuSFactoModelConfig(
|
||||
# # proposal network allows for significantly smaller sdf/color network
|
||||
# sdf_field=SDFFieldConfig(
|
||||
# use_grid_feature=True,
|
||||
# num_layers=2,
|
||||
# num_layers_color=2,
|
||||
# hidden_dim=256,
|
||||
# bias=0.5,
|
||||
# beta_init=0.8,
|
||||
# use_appearance_embedding=False,
|
||||
# ),
|
||||
# background_model="none",
|
||||
# eval_num_rays_per_chunk=2048,
|
||||
# ),
|
||||
# ),
|
||||
# optimizers={
|
||||
# "proposal_networks": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15),
|
||||
# "scheduler": MultiStepSchedulerConfig(max_steps=20001, milestones=(10000, 1500, 18000)),
|
||||
# },
|
||||
# "fields": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
# "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),
|
||||
# },
|
||||
# "field_background": {
|
||||
# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15),
|
||||
# "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001),
|
||||
# },
|
||||
# },
|
||||
# viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
|
||||
# vis="viewer",
|
||||
# )
|
||||
|
||||
|
||||
def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True):
|
||||
|
||||
@@ -43,7 +43,6 @@ from torch.nn import Parameter
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from typing_extensions import TypeVar
|
||||
|
||||
from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig
|
||||
from nerfstudio.cameras.cameras import CameraType
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from nerfstudio.configs.base_config import InstantiateConfig
|
||||
@@ -109,8 +108,6 @@ class DataManagerConfig(InstantiateConfig):
|
||||
"""Target class to instantiate."""
|
||||
data: Optional[Path] = None
|
||||
"""Source of data, may not be used by all models."""
|
||||
camera_optimizer: Optional[CameraOptimizerConfig] = None
|
||||
"""Specifies the camera pose optimizer used during training. Helpful if poses are noisy."""
|
||||
|
||||
|
||||
class DataManager(nn.Module):
|
||||
@@ -329,9 +326,6 @@ class VanillaDataManagerConfig(DataManagerConfig):
|
||||
new images. If -1, never pick new images."""
|
||||
eval_image_indices: Optional[Tuple[int, ...]] = (0,)
|
||||
"""Specifies the image indices to use during eval; if None, uses all."""
|
||||
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
|
||||
"""Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
|
||||
Record3D."""
|
||||
collate_fn: Callable[[Any], Any] = cast(Any, staticmethod(nerfstudio_collate))
|
||||
"""Specifies the collate function to use for the train and eval dataloaders."""
|
||||
camera_res_scale_factor: float = 1.0
|
||||
@@ -459,13 +453,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
|
||||
)
|
||||
self.iter_train_image_dataloader = iter(self.train_image_dataloader)
|
||||
self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch)
|
||||
self.train_camera_optimizer = self.config.camera_optimizer.setup(
|
||||
num_cameras=self.train_dataset.cameras.size, device=self.device
|
||||
)
|
||||
self.train_ray_generator = RayGenerator(
|
||||
self.train_dataset.cameras.to(self.device),
|
||||
self.train_camera_optimizer,
|
||||
)
|
||||
self.train_ray_generator = RayGenerator(self.train_dataset.cameras.to(self.device))
|
||||
|
||||
def setup_eval(self):
|
||||
"""Sets up the data loader for evaluation"""
|
||||
@@ -483,13 +471,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
|
||||
)
|
||||
self.iter_eval_image_dataloader = iter(self.eval_image_dataloader)
|
||||
self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch)
|
||||
self.eval_camera_optimizer = self.config.camera_optimizer.setup(
|
||||
num_cameras=self.eval_dataset.cameras.size, device=self.device
|
||||
)
|
||||
self.eval_ray_generator = RayGenerator(
|
||||
self.eval_dataset.cameras.to(self.device),
|
||||
self.eval_camera_optimizer,
|
||||
)
|
||||
self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device))
|
||||
# for loading full images
|
||||
self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader(
|
||||
input_dataset=self.eval_dataset,
|
||||
@@ -545,13 +527,4 @@ class VanillaDataManager(DataManager, Generic[TDataset]):
|
||||
Returns:
|
||||
A list of dictionaries containing the data manager's param groups.
|
||||
"""
|
||||
param_groups = {}
|
||||
|
||||
camera_opt_params = list(self.train_camera_optimizer.parameters())
|
||||
if self.config.camera_optimizer.mode != "off":
|
||||
assert len(camera_opt_params) > 0
|
||||
param_groups[self.config.camera_optimizer.param_group] = camera_opt_params
|
||||
else:
|
||||
assert len(camera_opt_params) == 0
|
||||
|
||||
return param_groups
|
||||
return {}
|
||||
@@ -211,13 +211,6 @@ class Trainer:
|
||||
"""
|
||||
optimizer_config = self.config.optimizers.copy()
|
||||
param_groups = self.pipeline.get_param_groups()
|
||||
camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer
|
||||
if camera_optimizer_config is not None and camera_optimizer_config.mode != "off":
|
||||
assert camera_optimizer_config.param_group not in optimizer_config
|
||||
optimizer_config[camera_optimizer_config.param_group] = {
|
||||
"optimizer": camera_optimizer_config.optimizer,
|
||||
"scheduler": camera_optimizer_config.scheduler,
|
||||
}
|
||||
return Optimizers(optimizer_config, param_groups)
|
||||
|
||||
def train(self) -> None:
|
||||
|
||||
@@ -18,7 +18,6 @@ Ray generator.
|
||||
from jaxtyping import Int
|
||||
from torch import Tensor, nn
|
||||
|
||||
from nerfstudio.cameras.camera_optimizers import CameraOptimizer
|
||||
from nerfstudio.cameras.cameras import Cameras
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
|
||||
@@ -34,10 +33,9 @@ class RayGenerator(nn.Module):
|
||||
|
||||
image_coords: Tensor
|
||||
|
||||
def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None:
|
||||
def __init__(self, cameras: Cameras) -> None:
|
||||
super().__init__()
|
||||
self.cameras = cameras
|
||||
self.pose_optimizer = pose_optimizer
|
||||
self.register_buffer("image_coords", cameras.get_image_coords(), persistent=False)
|
||||
|
||||
def forward(self, ray_indices: Int[Tensor, "num_rays 3"]) -> RayBundle:
|
||||
@@ -51,11 +49,8 @@ class RayGenerator(nn.Module):
|
||||
x = ray_indices[:, 2] # col indices
|
||||
coords = self.image_coords[y, x]
|
||||
|
||||
camera_opt_to_camera = self.pose_optimizer(c)
|
||||
|
||||
ray_bundle = self.cameras.generate_rays(
|
||||
camera_indices=c.unsqueeze(-1),
|
||||
coords=coords,
|
||||
camera_opt_to_camera=camera_opt_to_camera,
|
||||
)
|
||||
return ray_bundle
|
||||
@@ -48,6 +48,7 @@ from nerfstudio.model_components.scene_colliders import NearFarCollider
|
||||
from nerfstudio.model_components.shaders import NormalsShader
|
||||
from nerfstudio.models.base_model import Model, ModelConfig
|
||||
from nerfstudio.utils import colormaps
|
||||
from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig, CameraOptimizer
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -126,6 +127,8 @@ class NerfactoModelConfig(ModelConfig):
|
||||
"""Which implementation to use for the model."""
|
||||
appearance_embed_dim: int = 32
|
||||
"""Dimension of the appearance embedding."""
|
||||
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
|
||||
"""Config of the camera optimizer to use"""
|
||||
|
||||
|
||||
class NerfactoModel(Model):
|
||||
@@ -163,6 +166,9 @@ class NerfactoModel(Model):
|
||||
implementation=self.config.implementation,
|
||||
)
|
||||
|
||||
self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
|
||||
num_cameras=self.num_train_data, device="cpu" # for now initialize on cpu, these will be moved later
|
||||
)
|
||||
self.density_fns = []
|
||||
num_prop_nets = self.config.num_proposal_iterations
|
||||
# Build the proposal network(s)
|
||||
@@ -236,6 +242,12 @@ class NerfactoModel(Model):
|
||||
param_groups = {}
|
||||
param_groups["proposal_networks"] = list(self.proposal_networks.parameters())
|
||||
param_groups["fields"] = list(self.field.parameters())
|
||||
camera_opt_params = list(self.camera_optimizer.parameters())
|
||||
if self.config.camera_optimizer.mode != "off":
|
||||
assert len(camera_opt_params) > 0
|
||||
param_groups["camera_opt"] = camera_opt_params
|
||||
else:
|
||||
assert len(camera_opt_params) == 0
|
||||
return param_groups
|
||||
|
||||
def get_training_callbacks(
|
||||
@@ -273,6 +285,10 @@ class NerfactoModel(Model):
|
||||
return callbacks
|
||||
|
||||
def get_outputs(self, ray_bundle: RayBundle):
|
||||
# apply the camera optimizer pose tweaks
|
||||
if self.training:
|
||||
self.camera_optimizer.apply_to_raybundle(ray_bundle)
|
||||
|
||||
ray_samples: RaySamples
|
||||
ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
|
||||
field_outputs = self.field.forward(ray_samples, compute_normals=self.config.predict_normals)
|
||||
@@ -325,6 +341,10 @@ class NerfactoModel(Model):
|
||||
metrics_dict["psnr"] = self.psnr(outputs["rgb"], image)
|
||||
if self.training:
|
||||
metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"])
|
||||
|
||||
# Report the camera optimization metrics
|
||||
metrics_dict["camera_opt_translation"] = self.camera_optimizer.pose_adjustment[:, :3].norm()
|
||||
metrics_dict["camera_opt_rotation"] = self.camera_optimizer.pose_adjustment[:, 3:].norm()
|
||||
return metrics_dict
|
||||
|
||||
def get_loss_dict(self, outputs, batch, metrics_dict=None):
|
||||
|
||||
@@ -280,18 +280,6 @@ class VanillaPipeline(Pipeline):
|
||||
ray_bundle, batch = self.datamanager.next_train(step)
|
||||
model_outputs = self._model(ray_bundle) # train distributed data parallel model if world_size > 1
|
||||
metrics_dict = self.model.get_metrics_dict(model_outputs, batch)
|
||||
|
||||
if self.config.datamanager.camera_optimizer is not None:
|
||||
camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group
|
||||
if camera_opt_param_group in self.datamanager.get_param_groups():
|
||||
# Report the camera optimization metrics
|
||||
metrics_dict["camera_opt_translation"] = (
|
||||
self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm()
|
||||
)
|
||||
metrics_dict["camera_opt_rotation"] = (
|
||||
self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm()
|
||||
)
|
||||
|
||||
loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict)
|
||||
|
||||
return model_outputs, loss_dict, metrics_dict
|
||||
|
||||
Reference in new issue
Block a user