Merge pull request #1601 from ziyi-zhang/main

Radiance Field Loss / Many worlds

Author: Tom94

include/neural-graphics-primitives/common.h

@@ -38,6 +38,19 @@ using namespace tcnn;
 
 namespace ngp {
 
+// Training modes.
+// - NeRF: Standard volumetric reconstruction approach
+// - RFL (Radiance Field Loss): Promotes surface-like representations
+// - RFL-Relaxed: Hybrid approach that maintains NeRF-like volumetric properties while
+//   encouraging surface formation, resulting in faster rendering
+// For technical details, see: https://rgl.epfl.ch/publications/Zhang2025Radiance
+enum class ETrainMode : int {
+	Nerf,
+	Rfl,
+	RflRelax,
+};
+static constexpr const char* TrainModeStr = "Nerf\0Rfl\0RflRelax\0\0";
+
 enum class EMeshRenderMode : int {
 	Off,
 	VertexColors,

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

@@ -53,7 +53,9 @@ __global__ void render_nerf(
 	ENerfActivation density_activation,
 	ENerfActivation rgb_activation,
 	float min_transmittance,
-	bool train_in_linear_colors
+	bool train_in_linear_colors,
+	bool surface_rendering,
+	float surface_rendering_threshold
 ) {
 	uint32_t x = threadIdx.x + blockDim.x * blockIdx.x;
 	uint32_t y = threadIdx.y + blockDim.y * blockIdx.y;
@@ -142,14 +144,21 @@ __global__ void render_nerf(
 		// Composit color
 		float alpha = 1.f - __expf(-network_to_density(nerf_out.w, density_activation) * dt);
 		float weight = alpha * (1.0f - color.a);
-		color += vec4(network_to_rgb_vec(nerf_out.xyz(), rgb_activation) * weight, weight);
+		vec3 rgb = network_to_rgb_vec(nerf_out.xyz(), rgb_activation);
+		color += vec4(rgb * weight, weight);
 
 		if (weight > max_weight) {
 			max_weight = weight;
 			best_depth_candidate = lens.is_360() ? distance(pos, cam_pos) : dot(cam_fwd, pos - cam_pos);
 		}
 
-		if (color.a > (1.0f - min_transmittance)) {
+		if (surface_rendering && alpha >= surface_rendering_threshold) {
+			// Surface rendering: return the first surface point that has a sufficient occupancy
+			color.rgb() = rgb;
+			color.a = 1.0f;
+			best_depth_candidate = lens.is_360() ? distance(pos, cam_pos) : dot(cam_fwd, pos - cam_pos);
+			alive = false;
+		} else if (color.a > (1.0f - min_transmittance)) {
 			color /= color.a;
 			alive = false;
 		}

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

@@ -76,7 +76,11 @@ __global__ void train_nerf(
 	const vec3* __restrict__ exposure,
 	vec3* __restrict__ exposure_gradient,
 	float depth_supervision_lambda,
-	float near_distance
+	float near_distance,
+
+	uint32_t training_step,
+	ETrainMode training_mode,
+	uint32_t rfl_warmup_steps
 ) {
 	const uint32_t i = threadIdx.x + blockIdx.x * blockDim.x;
 
@@ -127,6 +131,39 @@ __global__ void train_nerf(
 	const float startt = advance_n_steps(tminmax.x, cone_angle, random_val(rng));
 	vec3 idir = vec3(1.0f) / ray.d;
 
+	if (train_with_random_bg_color) {
+		background_color = random_val_3d(rng);
+	}
+
+	vec3 pre_envmap_background_color = background_color = srgb_to_linear(background_color);
+
+	// Composit background behind envmap
+	vec4 envmap_value;
+	if (envmap) {
+		envmap_value = read_envmap(envmap, ray.d);
+		background_color = envmap_value.rgb() + background_color * (1.0f - envmap_value.a);
+	}
+
+	vec3 exposure_scale = exp(0.6931471805599453f * exposure[img]);
+
+	// Need rgbtarget before the first pass in RFL training mode
+	vec3 rgbtarget;
+	if (train_in_linear_colors || color_space == EColorSpace::Linear) {
+		rgbtarget = exposure_scale * texsamp.rgb() + (1.0f - texsamp.a) * background_color;
+
+		if (!train_in_linear_colors) {
+			rgbtarget = linear_to_srgb(rgbtarget);
+			background_color = linear_to_srgb(background_color);
+		}
+	} else if (color_space == EColorSpace::SRGB) {
+		background_color = linear_to_srgb(background_color);
+		if (texsamp.a > 0) {
+			rgbtarget = linear_to_srgb(exposure_scale * texsamp.rgb() / texsamp.a) * texsamp.a + (1.0f - texsamp.a) * background_color;
+		} else {
+			rgbtarget = background_color;
+		}
+	}
+
 	// first pass to compute an accurate number of steps
 	uint32_t j = 0;
 	float t = startt;
@@ -135,6 +172,7 @@ __global__ void train_nerf(
 
 	vec4 color = vec4(0.0f);
 	vec3 hitpoint = vec3(0.0f);
+	vec3 loss_bg = vec3(0.0f);
 
 	bool alive = valid;
 
@@ -180,8 +218,10 @@ __global__ void train_nerf(
 		// Composit color
 		float alpha = 1.f - __expf(-network_to_density(nerf_out.w, density_activation) * dt);
 		float weight = alpha * (1.0f - color.a);
-		color += vec4(network_to_rgb_vec(nerf_out.xyz(), rgb_activation) * weight, weight);
+		vec3 rgb = network_to_rgb_vec(nerf_out.rgb(), rgb_activation);
+		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()) {
@@ -209,40 +249,9 @@ __global__ void train_nerf(
 		numsteps_out[ray_idx*2+1] = base;
 	}
 
-	if (train_with_random_bg_color) {
-		background_color = random_val_3d(rng);
-	}
-
-	vec3 pre_envmap_background_color = background_color = srgb_to_linear(background_color);
-
-	// Composit background behind envmap
-	vec4 envmap_value;
-	if (envmap) {
-		envmap_value = read_envmap(envmap, ray.d);
-		background_color = envmap_value.rgb() + background_color * (1.0f - envmap_value.a);
-	}
-
-	vec3 exposure_scale = exp(0.6931471805599453f * exposure[img]);
-
-	vec3 rgbtarget;
-	if (train_in_linear_colors || color_space == EColorSpace::Linear) {
-		rgbtarget = exposure_scale * texsamp.rgb() + (1.0f - texsamp.a) * background_color;
-
-		if (!train_in_linear_colors) {
-			rgbtarget = linear_to_srgb(rgbtarget);
-			background_color = linear_to_srgb(background_color);
-		}
-	} else if (color_space == EColorSpace::SRGB) {
-		background_color = linear_to_srgb(background_color);
-		if (texsamp.a > 0) {
-			rgbtarget = linear_to_srgb(exposure_scale * texsamp.rgb() / texsamp.a) * texsamp.a + (1.0f - texsamp.a) * background_color;
-		} else {
-			rgbtarget = background_color;
-		}
-	}
-
 	if (1.0f - color.a >= EPSILON) {
 		color.rgb() += (1.0f - color.a) * background_color;
+		loss_bg += (1.0f - color.a) * loss_and_gradient(rgbtarget, background_color, loss_type).loss;
 	}
 
 	// Step again, this time computing loss
@@ -299,6 +308,7 @@ __global__ void train_nerf(
 
 	// now do it again computing gradients
 	vec4 color2 = vec4(0.0f);
+	vec3 loss_bg2 = vec3(0.0f);
 	float depth2 = 0.0f;
 	t = startt;
 	j = 0;
@@ -366,7 +376,42 @@ __global__ void train_nerf(
 
 		// 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();
-		const vec3 dloss_by_drgb = weight * lg.gradient;
+
+		float density_derivative = network_to_density_derivative(float(local_network_output[3]), density_activation);
+		const float depth_suffix = depth - depth2;
+		const float depth_supervision = depth_loss_gradient * (T * local_depth - depth_suffix);
+
+		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)
+			);
+		} 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.
+			const vec3 rgb_bg = suffix / fmaxf(1e-6f, T);
+			const vec3 rgb_lerp = (1 - alpha) * rgb_bg + alpha * rgb;
+			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)
+			);
+		} 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)
+			);
+		}
 
 		tvec<network_precision_t, 4> local_dL_doutput;
 
@@ -375,14 +420,6 @@ __global__ void train_nerf(
 		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]));
 
-		float density_derivative = network_to_density_derivative(float(local_network_output[3]), density_activation);
-		const float depth_suffix = depth - depth2;
-		const float depth_supervision = depth_loss_gradient * (T * local_depth - depth_suffix);
-
-		float dloss_by_dmlp = density_derivative * (
-			dt * (dot(lg.gradient, T * rgb - suffix) + depth_supervision)
-		);
-
 		//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;
 

include/neural-graphics-primitives/testbed.h

@@ -819,6 +819,9 @@ class Testbed {
 			default_rng_t density_grid_rng;
 			int view = 0;
 
+			ETrainMode train_mode = ETrainMode::RflRelax;
+			int rfl_warmup_steps = 1000;
+
 			float depth_supervision_lambda = 0.f;
 
 			GPUMemory<float> sharpness_grid;
@@ -880,6 +883,9 @@ class Testbed {
 
 		float cone_angle_constant = 1.f / 256.f;
 
+		bool surface_rendering = false;
+		float surface_rendering_threshold = 0.5f;
+
 		bool visualize_cameras = false;
 
 		float render_min_transmittance = 0.01f;

scripts/run.py

@@ -40,6 +40,9 @@ 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("--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.")
@@ -146,6 +149,17 @@ def get_scene(scene):
 		print("NeRF training ray near_distance ", args.near_distance)
 		testbed.nerf.training.near_distance = args.near_distance
 
+	if args.train_mode:
+		if args.train_mode.lower() == "nerf":
+			testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
+		elif args.train_mode.lower() == "rfl":
+			testbed.nerf.training.train_mode = ngp.TrainMode.Rfl
+		elif args.train_mode.lower() == "rfl_relax" or args.train_mode.lower() == "rflrelax":
+			testbed.nerf.training.train_mode = ngp.TrainMode.RflRelax
+		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")
 
@@ -167,6 +181,9 @@ def get_scene(scene):
 		# Match nerf paper behaviour and train on a fixed bg.
 		testbed.nerf.training.random_bg_color = False
 
+		# Ensure that the training mode is set to NeRF.
+		testbed.nerf.training.train_mode = ngp.TrainMode.Nerf
+
 	old_training_step = 0
 	n_steps = args.n_steps
 
@@ -176,6 +193,7 @@ def get_scene(scene):
 	if n_steps < 0 and (not args.load_snapshot or args.gui):
 		n_steps = 35000
 
+	original_train_mode = ngp.TrainMode(testbed.nerf.training.train_mode)
 	tqdm_last_update = 0
 	if n_steps > 0:
 		with tqdm(desc="Training", total=n_steps, unit="steps") as t:
@@ -194,6 +212,16 @@ def get_scene(scene):
 					old_training_step = 0
 					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
+
 				now = time.monotonic()
 				if now - tqdm_last_update > 0.1:
 					t.update(testbed.training_step - old_training_step)

src/nerf_loader.cu

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

src/python_api.cu

@@ -308,6 +308,12 @@ PYBIND11_MODULE(pyngp, m) {
 
 	m.def("free_temporary_memory", &free_all_gpu_memory_arenas);
 
+	py::enum_<ETrainMode>(m, "TrainMode")
+		.value("Nerf", ETrainMode::Nerf)
+		.value("Rfl", ETrainMode::Rfl)
+		.value("RflRelax", ETrainMode::RflRelax)
+		.export_values();
+
 	py::enum_<ETestbedMode>(m, "TestbedMode")
 		.value("Nerf", ETestbedMode::Nerf)
 		.value("Sdf", ETestbedMode::Sdf)
@@ -793,6 +799,8 @@ PYBIND11_MODULE(pyngp, m) {
 		//.def_readonly("focal_lengths", &Testbed::Nerf::Training::focal_lengths) // use training.dataset.metadata instead
 		.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)