From 767dc100d1800a7924a2343fb81431f8af2cfbaf Mon Sep 17 00:00:00 2001
From: Justin <justin.g.kerr@gmail.com>
Date: Tue, 20 Jun 2023 05:45:58 +0000
Subject: [PATCH] param

---
 nerfstudio/cameras/camera_optimizers.py | 10 ++++--
 nerfstudio/configs/method_configs.py    | 41 +++++++++++++++----------
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py
index 706e03a6e..edb774bee 100644
--- a/nerfstudio/cameras/camera_optimizers.py
+++ b/nerfstudio/cameras/camera_optimizers.py
@@ -28,8 +28,8 @@ from torch import Tensor, nn
 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.cameras.rays import RayBundle
+from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.utils import poses as pose_utils
 
 
@@ -42,12 +42,13 @@ 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
+    trans_l2_penalty: float = 1e-2
     """L2 penalty on translation parameters."""
 
     rot_l2_penalty: float = 1e-3
     """L2 penalty on rotation parameters."""
 
+
 class CameraOptimizer(nn.Module):
     """Layer that modifies camera poses to be optimized as well as the field during training."""
 
@@ -114,4 +115,7 @@ class CameraOptimizer(nn.Module):
         """
         Add a regularizer
         """
-        loss_dict['camera_opt_regularizer'] = self.pose_adjustment[:,:3].norm(dim=-1).mean() * self.config.trans_l2_penalty + self.pose_adjustment[:,3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
\ No newline at end of file
+        loss_dict["camera_opt_regularizer"] = (
+            self.pose_adjustment[:, :3].norm(dim=-1).mean() * self.config.trans_l2_penalty
+            + self.pose_adjustment[:, 3:].norm(dim=-1).mean() * self.config.rot_l2_penalty
+        )
diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py
index e25a78311..b553d3c4a 100644
--- a/nerfstudio/configs/method_configs.py
+++ b/nerfstudio/configs/method_configs.py
@@ -26,25 +26,32 @@ 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.dataparsers.blender_dataparser import BlenderDataParserConfig
+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.nerfstudio_dataparser import NerfstudioDataParserConfig
-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.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.sdfstudio_dataparser import \
+    SDFStudioDataParserConfig
+from nerfstudio.data.dataparsers.sitcoms3d_dataparser import \
+    Sitcoms3DDataParserConfig
 from nerfstudio.data.datasets.depth_dataset import DepthDataset
 from nerfstudio.data.datasets.sdf_dataset import SDFDataset
 from nerfstudio.data.datasets.semantic_dataset import SemanticDataset
-from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig
-from nerfstudio.engine.schedulers import (
-    CosineDecaySchedulerConfig,
-    ExponentialDecaySchedulerConfig,
-    MultiStepSchedulerConfig,
-)
+from nerfstudio.engine.optimizers import (AdamOptimizerConfig,
+                                          RAdamOptimizerConfig)
+from nerfstudio.engine.schedulers import (CosineDecaySchedulerConfig,
+                                          ExponentialDecaySchedulerConfig,
+                                          MultiStepSchedulerConfig)
 from nerfstudio.engine.trainer import TrainerConfig
-from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind
+from nerfstudio.field_components.temporal_distortions import \
+    TemporalDistortionKind
 from nerfstudio.fields.sdf_field import SDFFieldConfig
 from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig
 from nerfstudio.models.instant_ngp import InstantNGPModelConfig
@@ -103,12 +110,12 @@ method_configs["nerfacto"] = TrainerConfig(
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000),
         },
         "camera_opt": {
-            "optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8),
+            "optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
             "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000),
         },
     },
     viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
-    vis="viewer",
+    vis="viewer_beta",
 )
 
 # **For now I commented everything out since the datamanager config breaks **
@@ -198,7 +205,7 @@ method_configs["nerfacto-huge"] = TrainerConfig(
         "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",