Repository: NVlabs/instant-ngp
File: include/neural-graphics-primitives/fused_kernels/train_nerf.cuh (CUDA)

Code before:
		depth2 += weight * local_depth;
		float T = 1.0f - color2.a;

		t += dt;
		++j;

		// Compute and write loss stuff. Only needed if we want to actually output anything
		if (j > numsteps) {
			continue;
		}

		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;
		}

		// 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);
		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 < 1000) {
			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 * (

Review comment (bug):
This part *could* be moved into Python similar to what I describe in my comment below, but I think this depends on how necessary the NeRF warmup is.

Are there a couple of scenes where training fails completely if NeRF warmup is disabled? If so, let's leave it hardcoded here. But otherwise let's move it to Python alongside the rest of the training schedule.

Code after:
		float local_depth = distance(pos, ray.o);
		depth2 += weight * local_depth;
		float T = 1.0f - color2.a;

		t += dt;
		++j;

		// Compute and write loss stuff. Only needed if we want to actually output anything
		if (j > numsteps) {
			continue;
		}

		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;
		}

		// 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);
		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;