tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-12 19:04:03 +08:00
Code Activity

add l2 penalties on optimizers

Browse Source
This commit is contained in:
Justin committed 2023-06-18 06:09:46 +00:00
1 parent ae5ef15bed
commit be47ffca81
3 files changed
+63 -50
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nerfstudio/cameras/camera_optimizers.py
+12
View File
@@ -43,6 +43,11 @@ class CameraOptimizerConfig(InstantiateConfig):
mode: Literal["off", "SO3xR3", "SE3"] = "off"
"""Pose optimization strategy to use. If enabled, we recommend SO3xR3."""
trans_l2_penalty: float = 1e-3
"""L2 penalty on translation parameters."""
rot_l2_penalty: float = 1e-2
"""L2 penalty on rotation parameters."""
class CameraOptimizer(nn.Module):
"""Layer that modifies camera poses to be optimized as well as the field during training."""
@@ -105,3 +110,10 @@ class CameraOptimizer(nn.Module):
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()
def get_loss_dict(self, loss_dict: dict) -> dict:
"""
Add a regularizer
"""
loss_dict['camera_opt_regularizer'] = self.pose_adjustment[:,:3].norm(dim=-1).mean() * self.config.trans_l2_penalty
loss_dict['camera_rot_regularizer'] = self.pose_adjustment[:,3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
nerfstudio/configs/method_configs.py
+48 -49
View File
@@ -103,7 +103,7 @@ method_configs["nerfacto"] = TrainerConfig(
"scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
},
"camera_opt": {
"optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2),
"optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000),
},
},
@@ -156,54 +156,53 @@ method_configs["nerfacto"] = TrainerConfig(
# 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["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,
),
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),
},
"camera_opt": {
"optimizer": RAdamOptimizerConfig(lr=6e-4, eps=1e-8),
"scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000),
}
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
vis="viewer",
)
# method_configs["depth-nerfacto"] = TrainerConfig(
# method_name="depth-nerfacto",
nerfstudio/models/nerfacto.py
+3 -1
View File
@@ -127,7 +127,7 @@ class NerfactoModelConfig(ModelConfig):
"""Which implementation to use for the model."""
appearance_embed_dim: int = 32
"""Dimension of the appearance embedding."""
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3")
"""Config of the camera optimizer to use"""
@@ -367,6 +367,8 @@ class NerfactoModel(Model):
loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
outputs["rendered_pred_normal_loss"]
)
# Add losses from the camera optimizer (like L2 regularization)
self.camera_optimizer.get_loss_dict(loss_dict)
return loss_dict
def get_image_metrics_and_images(