feat(rfl): unify rfl and rflrelax training schedules as part of run.py

Author: Tom94

dependencies/tiny-cuda-nn

@@ -19,8 +19,7 @@
 
 using namespace ngp;
 
-__launch_bounds__(128, 4)
-__global__ void render_nerf(
+__launch_bounds__(128, 4) __global__ void render_nerf(
 	uint32_t sample_index,
 	ivec2 resolution,
 	vec2 focal_length,
@@ -65,7 +64,8 @@ __global__ void render_nerf(
 
 	vec2 pixel_offset = ld_random_pixel_offset(snap_to_pixel_centers ? 0 : sample_index);
 	vec2 uv = vec2{(float)x + pixel_offset.x, (float)y + pixel_offset.y} / vec2(resolution);
-	mat4x3 camera = get_xform_given_rolling_shutter({camera_matrix0, camera_matrix1}, rolling_shutter, uv, ld_random_val(sample_index, idx * 72239731));
+	mat4x3 camera =
+		get_xform_given_rolling_shutter({camera_matrix0, camera_matrix1}, rolling_shutter, uv, ld_random_val(sample_index, idx * 72239731));
 
 	Ray ray = uv_to_ray(
 		sample_index,
@@ -108,7 +108,9 @@ __global__ void render_nerf(
 		vec3 pos = cam_pos;
 
 		if (alive) {
-			t = if_unoccupied_advance_to_next_occupied_voxel(t, cone_angle, ray, idir, density_grid, min_mip, max_mip, render_aabb, render_aabb_to_local);
+			t = if_unoccupied_advance_to_next_occupied_voxel(
+				t, cone_angle, ray, idir, density_grid, min_mip, max_mip, render_aabb, render_aabb_to_local
+			);
 			alive &= t < MAX_DEPTH();
 			if (alive) {
 				pos = ray(t);

include/neural-graphics-primitives/fused_kernels/render_nerf.cuh

@@ -78,9 +78,7 @@ __global__ void train_nerf(
 	float depth_supervision_lambda,
 	float near_distance,
 
-	uint32_t training_step,
-	ETrainMode training_mode,
-	uint32_t rfl_warmup_steps
+	ETrainMode training_mode
 ) {
 	const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
 
@@ -222,6 +220,7 @@ __global__ void train_nerf(
 		color += vec4(rgb * weight, weight);
 
 		loss_bg += weight * loss_and_gradient(rgbtarget, rgb, loss_type).loss;
+
 		hitpoint += weight * pos;
 
 		if (1.0f - color.a < EPSILON || j >= NERF_STEPS()) {
@@ -233,7 +232,7 @@ __global__ void train_nerf(
 	hitpoint /= color.a;
 
 	uint32_t numsteps = j;
-	uint32_t base = atomicAdd(numsteps_counter, numsteps);	 // first entry in the array is a counter
+	uint32_t base = atomicAdd(numsteps_counter, numsteps); // first entry in the array is a counter
 	numsteps = min(max_samples - min(max_samples, base), numsteps);
 	bool can_write = numsteps > 0;
 
@@ -245,8 +244,8 @@ __global__ void train_nerf(
 	if (can_write) {
 		ray_idx = atomicAdd(ray_counter, 1);
 		ray_indices_out[ray_idx] = i;
-		numsteps_out[ray_idx*2+0] = numsteps;
-		numsteps_out[ray_idx*2+1] = base;
+		numsteps_out[ray_idx * 2 + 0] = numsteps;
+		numsteps_out[ray_idx * 2 + 1] = base;
 	}
 
 	if (1.0f - color.a >= EPSILON) {
@@ -258,7 +257,8 @@ __global__ void train_nerf(
 	LossAndGradient lg = loss_and_gradient(rgbtarget, color.rgb(), loss_type);
 	lg.loss /= img_pdf * uv_pdf;
 
-	float target_depth = ray_length * ((depth_supervision_lambda > 0.0f && metadata[img].depth) ? read_depth(uv, resolution, metadata[img].depth) : -1.0f);
+	float target_depth = ray_length *
+		((depth_supervision_lambda > 0.0f && metadata[img].depth) ? read_depth(uv, resolution, metadata[img].depth) : -1.0f);
 	LossAndGradient lg_depth = loss_and_gradient(vec3(target_depth), vec3(depth), depth_loss_type);
 	float depth_loss_gradient = target_depth > 0.0f ? depth_supervision_lambda * lg_depth.gradient.x : 0;
 
@@ -286,25 +286,28 @@ __global__ void train_nerf(
 
 		if (sharpness_data && aabb.contains(hitpoint)) {
 			ivec2 sharpness_pos = clamp(ivec2(uv * vec2(sharpness_resolution)), 0, sharpness_resolution - 1);
-			float sharp = sharpness_data[img * product(sharpness_resolution) + sharpness_pos.y * sharpness_resolution.x + sharpness_pos.x] + 1e-6f;
+			float sharp = sharpness_data[img * product(sharpness_resolution) + sharpness_pos.y * sharpness_resolution.x + sharpness_pos.x] +
+				1e-6f;
 
 			// The maximum value of positive floats interpreted in uint format is the same as the maximum value of the floats.
-			float grid_sharp = __uint_as_float(atomicMax((uint32_t*)&cascaded_grid_at(hitpoint, sharpness_grid, mip_from_pos(hitpoint, max_mip)), __float_as_uint(sharp)));
+			float grid_sharp = __uint_as_float(
+				atomicMax((uint32_t*)&cascaded_grid_at(hitpoint, sharpness_grid, mip_from_pos(hitpoint, max_mip)), __float_as_uint(sharp))
+			);
 			grid_sharp = fmaxf(sharp, grid_sharp); // atomicMax returns the old value, so compute the new one locally.
 
 			mean_loss *= fmaxf(sharp / grid_sharp, 0.01f);
 		}
 
-		deposit_val(idx.x,   idx.y,   (1 - weight.x) * (1 - weight.y) * mean_loss);
-		deposit_val(idx.x+1, idx.y,        weight.x  * (1 - weight.y) * mean_loss);
-		deposit_val(idx.x,   idx.y+1, (1 - weight.x) *      weight.y  * mean_loss);
-		deposit_val(idx.x+1, idx.y+1,      weight.x  *      weight.y  * mean_loss);
+		deposit_val(idx.x, idx.y, (1 - weight.x) * (1 - weight.y) * mean_loss);
+		deposit_val(idx.x + 1, idx.y, weight.x * (1 - weight.y) * mean_loss);
+		deposit_val(idx.x, idx.y + 1, (1 - weight.x) * weight.y * mean_loss);
+		deposit_val(idx.x + 1, idx.y + 1, weight.x * weight.y * mean_loss);
 	}
 
 	loss_scale /= n_rays;
 
 	const float output_l2_reg = rgb_activation == ENerfActivation::Exponential ? 1e-4f : 0.0f;
-	const float output_l1_reg_density = 0.0f;// *mean_density_ptr < NERF_MIN_OPTICAL_THICKNESS() ? 1e-4f : 0.0f;
+	const float output_l1_reg_density = 0.0f; // *mean_density_ptr < NERF_MIN_OPTICAL_THICKNESS() ? 1e-4f : 0.0f;
 
 	// now do it again computing gradients
 	vec4 color2 = vec4(0.0f);
@@ -369,12 +372,13 @@ __global__ void train_nerf(
 			continue;
 		}
 
-		coords_out(j-1)->copy(*(NerfCoordinate*)&nerf_in[0], coords_out.stride_in_bytes);
+		coords_out(j - 1)->copy(*(NerfCoordinate*)&nerf_in[0], coords_out.stride_in_bytes);
 		if (max_level_rand_training) {
-			max_level_ptr[j-1] = max_level;
+			max_level_ptr[j - 1] = max_level;
 		}
 
-		// we know the suffix of this ray compared to where we are up to. note the suffix depends on this step's alpha as suffix = (1-alpha)*(somecolor), so dsuffix/dalpha = -somecolor = -suffix/(1-alpha)
+		// we know the suffix of this ray compared to where we are up to. note the suffix depends on this step's alpha as suffix =
+		// (1-alpha)*(somecolor), so dsuffix/dalpha = -somecolor = -suffix/(1-alpha)
 		const vec3 suffix = color.rgb() - color2.rgb();
 
 		float density_derivative = network_to_density_derivative(float(local_network_output[3]), density_activation);
@@ -383,17 +387,13 @@ __global__ void train_nerf(
 
 		vec3 dloss_by_drgb;
 		float dloss_by_dmlp;
-		if (training_mode == ETrainMode::Rfl && training_step < rfl_warmup_steps) {
-			training_mode = ETrainMode::Nerf; // Warm up training
-		}
+
 		if (training_mode == ETrainMode::Rfl) {
 			// Radiance field loss
 			LossAndGradient local_lg = loss_and_gradient(rgbtarget, rgb, loss_type);
 			loss_bg2 += weight * local_lg.loss;
 			dloss_by_drgb = weight * local_lg.gradient;
-			dloss_by_dmlp = density_derivative * (
-				dt * sum(T * local_lg.loss - (loss_bg - loss_bg2) + depth_supervision)
-			);
+			dloss_by_dmlp = density_derivative * (dt * sum(T * local_lg.loss - (loss_bg - loss_bg2) + depth_supervision));
 		} else if (training_mode == ETrainMode::RflRelax) {
 			// In-between volume reconstruction and surface reconstruction.
 			// This is different from the relaxation in the paper, but is much simpler and also promotes surfaces.
@@ -402,32 +402,33 @@ __global__ void train_nerf(
 			LossAndGradient local_lg = loss_and_gradient(rgbtarget, rgb_lerp, loss_type);
 
 			dloss_by_drgb = weight * local_lg.gradient;
-			dloss_by_dmlp = density_derivative * (
-				dt * (dot(local_lg.gradient, T * rgb - suffix) + depth_supervision)
-			);
+			dloss_by_dmlp = density_derivative * (dt * (dot(local_lg.gradient, T * rgb - suffix) + depth_supervision));
 		} else {
 			// The original NeRF loss
 			dloss_by_drgb = weight * lg.gradient;
-			dloss_by_dmlp = density_derivative * (
-				dt * (dot(lg.gradient, T * rgb - suffix) + depth_supervision)
-			);
+			dloss_by_dmlp = density_derivative * (dt * (dot(lg.gradient, T * rgb - suffix) + depth_supervision));
 		}
 
 		tvec<network_precision_t, 4> local_dL_doutput;
 
 		// chain rule to go from dloss/drgb to dloss/dmlp_output
-		local_dL_doutput[0] = loss_scale * (dloss_by_drgb.x * network_to_rgb_derivative(local_network_output[0], rgb_activation) + fmaxf(0.0f, output_l2_reg * (float)local_network_output[0])); // Penalize way too large color values
-		local_dL_doutput[1] = loss_scale * (dloss_by_drgb.y * network_to_rgb_derivative(local_network_output[1], rgb_activation) + fmaxf(0.0f, output_l2_reg * (float)local_network_output[1]));
-		local_dL_doutput[2] = loss_scale * (dloss_by_drgb.z * network_to_rgb_derivative(local_network_output[2], rgb_activation) + fmaxf(0.0f, output_l2_reg * (float)local_network_output[2]));
-
-		//static constexpr float mask_supervision_strength = 1.f; // we are already 'leaking' mask information into the nerf via the random bg colors; setting this to eg between 1 and  100 encourages density towards 0 in such regions.
-		//dloss_by_dmlp += (texsamp.a<0.001f) ? mask_supervision_strength * weight : 0.f;
-
-		local_dL_doutput[3] =
-			loss_scale * dloss_by_dmlp +
-			(float(local_network_output[3]) < 0.0f ? -output_l1_reg_density : 0.0f) +
+		local_dL_doutput[0] = loss_scale *
+			(dloss_by_drgb.x * network_to_rgb_derivative(local_network_output[0], rgb_activation) +
+			 fmaxf(0.0f, output_l2_reg * (float)local_network_output[0])); // Penalize way too large color values
+		local_dL_doutput[1] = loss_scale *
+			(dloss_by_drgb.y * network_to_rgb_derivative(local_network_output[1], rgb_activation) +
+			 fmaxf(0.0f, output_l2_reg * (float)local_network_output[1]));
+		local_dL_doutput[2] = loss_scale *
+			(dloss_by_drgb.z * network_to_rgb_derivative(local_network_output[2], rgb_activation) +
+			 fmaxf(0.0f, output_l2_reg * (float)local_network_output[2]));
+
+		// static constexpr float mask_supervision_strength = 1.f; // we are already 'leaking' mask information into the nerf via the random
+		// bg colors; setting this to eg between 1 and  100 encourages density towards 0 in such regions. dloss_by_dmlp +=
+		// (texsamp.a<0.001f) ? mask_supervision_strength * weight : 0.f;
+
+		local_dL_doutput[3] = loss_scale * dloss_by_dmlp + (float(local_network_output[3]) < 0.0f ? -output_l1_reg_density : 0.0f) +
 			(float(local_network_output[3]) > -10.0f && local_depth < near_distance ? 1e-4f : 0.0f);
-			;
+		;
 
 		*(tvec<network_precision_t, 4>*)dloss_doutput = local_dL_doutput;
 		dloss_doutput += padded_output_width;

include/neural-graphics-primitives/fused_kernels/train_nerf.cuh

@@ -820,7 +820,6 @@ class Testbed {
 			int view = 0;
 
 			ETrainMode train_mode = ETrainMode::RflRelax;
-			int rfl_warmup_steps = 1000;
 
 			float depth_supervision_lambda = 0.f;
 

include/neural-graphics-primitives/testbed.h

@@ -40,9 +40,11 @@ def parse_args():
 	parser.add_argument("--test_transforms", default="", help="Path to a nerf style transforms json from which we will compute PSNR.")
 	parser.add_argument("--near_distance", default=-1, type=float, help="Set the distance from the camera at which training rays start for nerf. <0 means use ngp default")
 	parser.add_argument("--exposure", default=0.0, type=float, help="Controls the brightness of the image. Positive numbers increase brightness, negative numbers decrease it.")
+
 	parser.add_argument("--train_mode", default="", type=str, help="The training mode to use. Can be 'nerf', 'rfl', 'rfl_relax'. If not specified, the default mode will be used.")
 	parser.add_argument("--rfl_warmup_steps", type=int, default=1000, help="Number of steps to train in NeRF mode before switching to RFL mode. Default is 1000. Only used if --train_mode is set to 'rfl'.")
-	parser.add_argument("--no_rflrelax_training_schedule", action="store_true", help="Disable RFL training schedule for RflRelax mode (active between steps 15k-30k).")
+	parser.add_argument("--rflrelax_begin_step", type=int, default=15000, help="First training step in which RflRelax mode is used. Default is 15000. Only used if --train_mode is set to 'rflrelax'.")
+	parser.add_argument("--rflrelax_end_step", type=int, default=30000, help="Last training step in which RflRelax mode is used. Default is 30000. Only used if --train_mode is set to 'rflrelax'.")
 
 	parser.add_argument("--screenshot_transforms", default="", help="Path to a nerf style transforms.json from which to save screenshots.")
 	parser.add_argument("--screenshot_frames", nargs="*", help="Which frame(s) to take screenshots of.")
@@ -159,8 +161,6 @@ def get_scene(scene):
 		else:
 			raise ValueError(f"Unknown train mode: {args.train_mode}")
 
-	testbed.nerf.training.rfl_warmup_steps = args.rfl_warmup_steps
-
 	if args.nerf_compatibility:
 		print(f"NeRF compatibility mode enabled")
 
@@ -183,8 +183,11 @@ def get_scene(scene):
 		testbed.nerf.training.random_bg_color = False
 
 		# Ensure that the training mode is set to NeRF.
+		if testbed.nerf.training.train_mode != ngp.TrainMode.Nerf:
+			print(f"Warning: forcing train mode to NeRF for nerf compatibility (was {testbed.nerf.training.train_mode})")
 		testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
 
+
 	old_training_step = 0
 	n_steps = args.n_steps
 
@@ -196,14 +199,15 @@ def get_scene(scene):
 
 	original_train_mode = ngp.TrainMode(testbed.nerf.training.train_mode)
 	prev_train_mode = original_train_mode
+	use_training_schedule = True
 
 	tqdm_last_update = 0
 	if n_steps > 0:
 		with tqdm(desc="Training", total=n_steps, unit="steps") as t:
 			while testbed.frame():
-				if prev_train_mode != testbed.nerf.training.train_mode and not args.no_rflrelax_training_schedule:
-					print("Disabling RflRelax training schedule due to UI train mode change")
-					args.no_rflrelax_training_schedule = True
+				if prev_train_mode != testbed.nerf.training.train_mode and use_training_schedule:
+					print("Disabling Rfl/RflRelax training schedule due to UI train mode change")
+					use_training_schedule = False
 
 				if testbed.want_repl():
 					repl(testbed)
@@ -221,13 +225,21 @@ def get_scene(scene):
 					t.reset()
 
 				# Rfl-relax training schedule
-				progress_fraction = float(testbed.training_step) / n_steps
-				if original_train_mode == ngp.TrainMode.RflRelax and not args.no_rflrelax_training_schedule:
-					# By default only enable RflRelax mode between 15k and 30k steps
-					if 3/7 <= progress_fraction < 6/7:
-						testbed.nerf.training.train_mode = ngp.TrainMode.RflRelax
-					else:
-						testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
+				if use_training_schedule:
+					if original_train_mode == ngp.TrainMode.RflRelax:
+						# By default only enable RflRelax mode in the middle of training. Start with NeRF mode,
+						# then switch to RflRelax mode to "sueface-ify" the scene, then switch back to NeRF mode
+						# at the very and for fine tuning.
+						if args.rflrelax_begin_step <= testbed.training_step < args.rflrelax_end_step:
+							testbed.nerf.training.train_mode = ngp.TrainMode.RflRelax
+						else:
+							testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
+					elif original_train_mode == ngp.TrainMode.Rfl:
+						# Start in NeRF mode, then switch to RFL mode after a warmup period
+						if testbed.training_step > args.rfl_warmup_steps:
+							testbed.nerf.training.train_mode = ngp.TrainMode.Rfl
+						else:
+							testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
 
 				now = time.monotonic()
 				if now - tqdm_last_update > 0.1:

scripts/run.py

@@ -442,8 +442,8 @@ NerfDataset load_nerf(const std::vector<fs::path>& jsonpaths, float sharpen_amou
 		}
 
 		if (json.contains("from_mitsuba")) {
-  		    result.from_mitsuba = bool(json["from_mitsuba"]);
-  		}
+			result.from_mitsuba = bool(json["from_mitsuba"]);
+		}
 
 		if (json.contains("fix_premult")) {
 			fix_premult = (bool)json["fix_premult"];

src/nerf_loader.cu

@@ -800,7 +800,6 @@ PYBIND11_MODULE(pyngp, m) {
 		.def_readwrite("near_distance", &Testbed::Nerf::Training::near_distance)
 		.def_readwrite("density_grid_decay", &Testbed::Nerf::Training::density_grid_decay)
 		.def_readwrite("train_mode", &Testbed::Nerf::Training::train_mode)
-		.def_readwrite("rfl_warmup_steps", &Testbed::Nerf::Training::rfl_warmup_steps)
 		.def_readwrite("extrinsic_l2_reg", &Testbed::Nerf::Training::extrinsic_l2_reg)
 		.def_readwrite("extrinsic_learning_rate", &Testbed::Nerf::Training::extrinsic_learning_rate)
 		.def_readwrite("intrinsic_l2_reg", &Testbed::Nerf::Training::intrinsic_l2_reg)