nerfstudio icon indicating copy to clipboard operation
nerfstudio copied to clipboard

use gsplat strategy interface to simplify splatfacto

Open jb-ye opened this issue 1 year ago • 5 comments

Given the latest updates of gsplat strategy interface, it is of interest to use the feature to benefit from multiple strategy implementations provided. The first step is to refactor existing splatfacto code meanwhile maintain the performance of current splafacto-big.

I spent some time debugging current gsplat implementation and found that we need to first merge this fix in order to reproduce the performance of current splatfacto-big. We don't necessarily need to merge the fix, this is just for producing baselines.

To benchmark simple_trainer.py in splatfacto configuration:

# using latest commit of gsplat and run in splatfacto-big equivalent configurations
python simple_trainer.py default --data_dir bicycle/ \
   --strategy.absgrad \
   --strategy.grow_grad2d 0.0006 \
   --strategy.prune_scale3d 0.5 \
   --strategy.refine_scale2d_stop_iter 4000 \
   --strategy.pause_refine_after_reset 300

The above setting in simple_trainer is "almost" equivalent to splatfacto-big (post this PR merge) with two differences:

  • splatfacto-big have coarse-to-fine training enabled by default. Without coarse-to-fine training, absgrad tends to create a lot more gaussians in early iterations, which might not be necessarily useful.
  • scene_scale calculation in nerfstudio colmap parser use abs(cam_positions) instead of 1.1 * norm(cam_positions). This uniformly affects all--strategy.*3d parameters.

Mipnerf360 bicycle

  • Baseline A (original Inria):

    • 7k iterations: PSNR=23.59, SSIM=0.662, GS=3.57M
    • 30k iterations: PSNR=25.19, SSIM=0.763, GS=6.06M
  • Baseline B (splatfacto-big in main):

    • 30k iterations: PSNR=25.25, SSIM=0.781, GS=5.36M
  • simple_trainer (using splatfacto configuration) + fix:

    • 7k iterations: PSNR=23.96, SSIM=0.714, GS=5.06M
    • 30k iterations: PSNR=25.22, SSIM=0.780, GS=6.46M
  • splatfacto-big (this PR) + fix:

    • 7k iterations: GS=3.9M
    • 30k iterations: PSNR=25.28, SSIM=0.781, GS=6.03M
  • splatfacto-big (this PR, densify_grad_thresh=0.0006):

    • 7k iterations: GS=3.0M
    • 30k iterations: PSNR=25.24, SSIM=0.777, GS=4.70M
  • splatfacto-big (this PR, densify_grad_thresh=0.0005):

    • 7k iterations: GS=3.8M
    • 30k iterations: PSNR=25.26, SSIM=0.781, GS=6.00M

Note that the default strategy is still a strong candidate if one chooses a proper configuration. The original Inria configuration is suboptimal.

jb-ye avatar Aug 22 '24 18:08 jb-ye

Looks great! thanks for taking this on. Are you planning on adding support for MCMC? it would be nice to remove hyperparameters from SplatfactoConfig and replace them with a sort of StrategyConfig which can be swapped out for different strategies. This might require a gsplat PR to format them in a compatible way with tyro.

I see that https://github.com/nerfstudio-project/gsplat/pull/358 was closed without merging, is that change reflected inside the splatfacto hyperparams for gradient threhsolds now?

kerrj avatar Aug 25 '24 22:08 kerrj

Are you planning on adding support for MCMC?

Yes, but I am also open if someone else wish to take on this task. This PR is part of my own learning about the new strategy interface. This is indeed very convenient.

it would be nice to remove hyperparameters from SplatfactoConfig and replace them with a sort of StrategyConfig which can be swapped out for different strategies. This might require a gsplat PR to format them in a compatible way with tyro.

Yes, it takes some thoughts about how to change the SplatfactoConfig.

I see that https://github.com/nerfstudio-project/gsplat/pull/358 was closed without merging, is that change reflected inside the splatfacto hyperparams for gradient threhsolds now?

Yes, I retuned the parameter, I believe they are now very close.

jb-ye avatar Aug 26 '24 15:08 jb-ye

Since this PR removed the --continue_cull_post_densification option, we will need to filter gaussians whose opacities are below 1/255. upon exporting.

jb-ye avatar Aug 26 '24 21:08 jb-ye

It seems num_sh_bases is no longer in the most recent version gsplat, we can fix this import error by manually defining it like this:

def num_sh_bases(degree: int) -> int:
    """
    Returns the number of spherical harmonic bases for a given degree.
    """
    return (degree + 1) ** 2

kerrj avatar Sep 04 '24 22:09 kerrj

@kerrj Did we add a new test for building nerfstudio dockers? Not sure how to make it pass.

jb-ye avatar Sep 05 '24 20:09 jb-ye

@jkulhanek for Docker action, any idea?

brentyi avatar Sep 11 '24 05:09 brentyi

I think the error is that the docker build ran from a forked repo, but fork doesn’t have permission to run it. We should only run the docker action on the main repo - disable running it on forked repos I think.

jkulhanek avatar Sep 11 '24 06:09 jkulhanek

I've added a PR which should fix it.

jkulhanek avatar Sep 12 '24 07:09 jkulhanek

@brentyi @kerrj I think this PR is good to go.

jb-ye avatar Sep 12 '24 15:09 jb-ye

PR looks reasonable to me!

As a thought for supporting MCMC the easiest way I can think of for adding this is just putting into the config:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: DefaultStrategy = dataclasses.field(default_factory=DefaultStrategy)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

Or, if preventing breaking changes is more important than CLI flag organization:

strategy_mode: Literal["default", "mcmc"] = "default"
default_strategy: Annotated[DefaultStrategy, tyro.conf.arg(name="")] = dataclasses.field(
    default_factory=DefaultStrategy
)
mcmc_strategy: MCMCStrategy = dataclasses.field(MCMCStrategy)

and then in the splatfacto constructor:

self.strategy = {
    "default": self.config.default_strategy,
    "mcmc": self.config.mcmc_strategy,
}[self.config.strategy_mode]

Hello,

Thank you for your awesome work.

I assume MCMC is not yet implemented. Do you have any update for when we could try this with nerfstudio?

abrahamezzeddine avatar Sep 20 '24 16:09 abrahamezzeddine

@abrahamezzeddine someone created a PR for mcmc. But I haven't gotten time to look at it.

jb-ye avatar Sep 20 '24 20:09 jb-ye

Hi, in this line, scene_scale is hard-coded to 1 https://github.com/nerfstudio-project/nerfstudio/blob/516fd7c9ac73e28db6a522df57d7b22e7b0d6756/nerfstudio/models/splatfacto.py#L332 From my understanding, scene_scale is the maximum of distance between average camera position and all camera positions. When I trace the code of nerfstudio's colmap dataparser, the default setting is to auto scale all the camera poses and normalized the camera positions to -1 and 1. This results in scene_scale = 1. However, in my application, I don't want to scale camera positions to -1 and 1. This results in scene_scale does not always = 1.

Sunnyhong0326 avatar Sep 23 '24 04:09 Sunnyhong0326

@Sunnyhong0326 this is a known limitation. Ideally, we should pass scene scale from parser to the model in nerfstudio framework. Let me see what I can do here

jb-ye avatar Sep 23 '24 17:09 jb-ye