Base extrinsics optimization LR decay on the respective optimizer's s…

…tep count rather than the global one Gradually growing datasets thus allow extrinsics learning of new images while old ones stay put
NVlabs · Mar 3, 2022 · 7b37eee · 7b37eee
1 parent af9434e
commit 7b37eee
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/include/neural-graphics-primitives/adam_optimizer.h b/include/neural-graphics-primitives/adam_optimizer.h
@@ -44,6 +44,10 @@ class AdamOptimizer {
 		m_state.variable -= actual_learning_rate * m_state.first_moment.cwiseQuotient(m_state.second_moment.cwiseSqrt() + T::Constant(m_hparams.epsilon));
 	}
 
+	uint32_t step() const {
+		return m_state.iter;
+	}
+
 	void set_learning_rate(float lr) {
 		m_hparams.learning_rate = lr;
 	}
@@ -62,7 +66,7 @@ class AdamOptimizer {
 
 	private:
 	struct State {
-		int iter = 0;
+		uint32_t iter = 0;
 		T first_moment = T::Zero();
 		T second_moment = T::Zero();
 		T variable = T::Zero();
@@ -76,8 +80,6 @@ class AdamOptimizer {
 	} m_hparams;
 };
 
-
-
 class RotationAdamOptimizer {
 public:
 	RotationAdamOptimizer(float learning_rate, float epsilon = 1e-08f, float beta1 = 0.9f, float beta2 = 0.99f) {
@@ -112,6 +114,10 @@ class RotationAdamOptimizer {
 		m_state.variable = result.axis() * result.angle();
 	}
 
+	uint32_t step() const {
+		return m_state.iter;
+	}
+
 	void set_learning_rate(float lr) {
 		m_hparams.learning_rate = lr;
 	}
@@ -126,7 +132,7 @@ class RotationAdamOptimizer {
 
 	private:
 	struct State {
-		int iter = 0;
+		uint32_t iter = 0;
 		Eigen::Vector3f first_moment = Eigen::Vector3f::Zero();
 		Eigen::Vector3f second_moment = Eigen::Vector3f::Zero();
 		Eigen::Vector3f variable = Eigen::Vector3f::Zero();

diff --git a/src/testbed_nerf.cu b/src/testbed_nerf.cu
@@ -2573,8 +2573,8 @@ void Testbed::train_nerf(uint32_t target_batch_size, uint32_t n_training_steps,
 				pos_gradient += m_nerf.training.cam_pos_offset[i].variable() * l2_reg;
 				rot_gradient += m_nerf.training.cam_rot_offset[i].variable() * l2_reg;
 
-				m_nerf.training.cam_pos_offset[i].set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_training_step / 2048)), m_optimizer->learning_rate()/1000.0f));
-				m_nerf.training.cam_rot_offset[i].set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_training_step / 2048)), m_optimizer->learning_rate()/1000.0f));
+				m_nerf.training.cam_pos_offset[i].set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_nerf.training.cam_pos_offset[i].step() / 128)), m_optimizer->learning_rate()/1000.0f));
+				m_nerf.training.cam_rot_offset[i].set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_nerf.training.cam_rot_offset[i].step() / 128)), m_optimizer->learning_rate()/1000.0f));
 
 				m_nerf.training.cam_pos_offset[i].step(pos_gradient);
 				m_nerf.training.cam_rot_offset[i].step(rot_gradient);
@@ -2598,7 +2598,7 @@ void Testbed::train_nerf(uint32_t target_batch_size, uint32_t n_training_steps,
 			Vector2f focal_length_gradient = m_nerf.training.cam_focal_length_gradient * per_camera_loss_scale;
 			float l2_reg = m_nerf.training.intrinsic_l2_reg;
 			focal_length_gradient += m_nerf.training.cam_focal_length_offset.variable() * l2_reg;
-			m_nerf.training.cam_focal_length_offset.set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_training_step / 2048)),m_optimizer->learning_rate() / 1000.0f));
+			m_nerf.training.cam_focal_length_offset.set_learning_rate(std::max(1e-3f * std::pow(0.33f, (float)(m_nerf.training.cam_focal_length_offset.step() / 128)),m_optimizer->learning_rate() / 1000.0f));
 			m_nerf.training.cam_focal_length_offset.step(focal_length_gradient);
 			m_nerf.training.update_metadata();
 		}