Implement gradient scaling

[jkulhanek]

This PR implements gradient scaling from Radiance Field Gradient Scaling for Unbiased Near-Camera Training

GradientScaling is added to the nerfacto model and enabled using flag --pipeline.model.use_gradient_scaling True

[jkulhanek]

I am not sure it helps:

Here are the full runs, please take a look:
https://wandb.ai/kulhanek/nerfstudio-baselines?workspace=user-kulhanek

[tancik]

Black was updated, use pip install -e .[dev] for the latest version.

[tancik]

Add example of how to call in docstring. More info here

[tancik]

Can we type and add dosctrings? Or does it break since we are overwriting?

[jkulhanek]

Typing and doctrings are not visible anywhere. I have added a wrapper function to provide the docstrings and typing support.

[salykova]

@jkulhanek @tancik

In the paper they set near_plane to 0 in all experiments. Does it make sense to do the same here?

[tancik]

@jkulhanek @tancik

In the paper they set near_plane to 0 in all experiments. Does it make sense to do the same here?

The near plane can be set with --pipeline.model.near-plane
If we set the gradient scaling on by default, then it makes sense to set the near plane to 0 by default. I'm fine with making this default, but we should test it on a few more scenes first to make sure that it better or equal to not using it.

[ichsan2895]

Can't wait to be implemented in the next version of Nerfstudio 💯

[ichsan2895]

I am not sure it helps:

Here are the full runs, please take a look: https://wandb.ai/kulhanek/nerfstudio-baselines?workspace=user-kulhanek

As stated in the paper, the PSNR will slightly change (not too much), but it more better qualitatively because less of floaters

[tancik]

LGTM. We should try experiment with it over the next few days. If we see better results on average, let's enable it by default.

[ichsan2895]

Sorry why I can't use it?

another attempt

[ichsan2895]

GradientScaling is added to the nerfacto model and enabled using flag --use-gradient-scaling=True

This seems works with --pipeline.model.use_gradient_scaling True instead of --use-gradient-scaling=True

Here is the show. I decreased the value of num-ray-per-chunk & others parameter to avoid CUDA OOM

Without Gradient Scale, CHKPT 48 K, num-rays-per-chunk 16384, num-rays-per-batch 2048 , num-rays-per-batch 2048

With Gradient Scale, CHKPT 48 K, num-rays-per-chunk 16384, num-rays-per-batch 2048 , num-rays-per-batch 2048

[jkulhanek]

Thanks for testing it on your scene! I guess we can optimize it to fit into the memory.

[jkulhanek]

Does anyone know if gradients in autograd.function can me modified in-place?

[ichsan2895]

Thanks for testing it on your scene! I guess we can optimize it to fit into the memory.

Both of tests will OOM CUDA, so I think the problem of OOM is not caused by grad scaling. The problem of OOM CUDA come from my own potato GPU which it has low VRAM

[jkulhanek]

Ok, I guess I will keep the current implementation then.

[ichsan2895]

Another testing come here...

Without Gradient Scaling

without.GradScal.mp4

With Gradient Scaling

with.GradScal.mp4