Abstracting out the GS heuristics into a Strategy class #278

liruilong940607 · 2024-07-11T03:12:20Z

The idea is to move the densification heuristics from the example scripts into gsplat library and package them into a class, so that the downstream users could simply use them via from gsplat.strategy import DefaultStrategy, MCMCStrategy.

The API is fairly simple:

params: torch.nn.ParameterDict = ...
optimizers: Dict[str, torch.optim.Optimizer] = ...
strategy = DefaultStrategy(<arguments>)
state = strategy.initialize_state()

for step in range(max_steps):
    render_color, render_alpha, info = rasterization(...)
    strategy.step_pre_backward(params, optimizers, state, step=step, info=info, <other arguments>)
    loss.backward()
    strategy.step_post_backward(params, optimizers, state, step=step, info=info, <other arguments>)

This would also allow users to switch between different strategies, and would make it more easily to embrace more GS heuristics. This is a joint efforts with @rtabrizi.

I also verified that by switch to DefaultStrategy and MCMCStrategy in the example script, nothing is affected. Now these two scripts have very little differences so we could potentially merge them into one script. But I'll leave it to another PR to minimize the code change of this PR

python simple_trainer.py

# Before this PR
> Step:  6999 {'mem': 7.582746505737305, 'ellipse_time': 424.46320390701294, 'num_GS': 4488772}
> {"psnr": 26.358938217163086, "ssim": 0.8353559374809265, "lpips": 0.12293213605880737, "ellipse_time": 0.026054859161376953, "num_GS": 4488772}

# After this PR
> Step:  6999 {'mem': 7.579144477844238, 'ellipse_time': 423.3516597747803, 'num_GS': 4481712}
> {"psnr": 26.348445892333984, "ssim": 0.8354260325431824, "lpips": 0.12293919920921326, "ellipse_time": 0.026077985763549805, "num_GS": 4481712}

python simple_trainer_mcmc.py 

# Before this PR
> Step:  6999 {'mem': 1.7906079292297363, 'ellipse_time': 247.97816705703735, 'num_GS': 1000000}
> {"psnr": 26.010623931884766, "ssim": 0.8161215782165527, "lpips": 0.15706905722618103, "ellipse_time": 0.009712855021158854, "num_GS": 1000000}

# After this PR
> Step:  6999 {'mem': 1.7903776168823242, 'ellipse_time': 239.95758247375488, 'num_GS': 1000000}
> {"psnr": 26.001811981201172, "ssim": 0.8158013224601746, "lpips": 0.1573895812034607, "ellipse_time": 0.009663820266723633, "num_GS": 1000000}

TODO:

Documentation. @liruilong940607
See if it fits with Splatfacto in nerfstudio. branch in nerfstudio

tentative doc preview: https://669107b6a724ef1f2c186e67--inquisitive-speculoos-27b42b.netlify.app/apis/strategy

* initial commit, rough design of modules * moved strategy.py into gsplat dir * latest commit for 1:1 meeting * removed ops, losses, and strategies now inherit from an abstract 'Strategy' class * GSs init handled in private method

…eterDict with newly-added 'activations' dict

…ject/gsplat into rtabrizi/refactor

…ss for revised ADS heuristics

…ice argument passing

kerrj · 2024-07-11T06:50:31Z

gsplat/strategy.py

+    def __init__(
+        self,
+        params: torch.nn.ParameterDict,
+        optimizers: List[torch.optim.Optimizer],


wouldn't this need to be a dict matching the key in the paramdict to the optimizer?

A single optimizer is enough as long as the param group has a "name" field defined. The reason I'm putting a List here is in case user wants to use different optimizer for different field. E.g., SGD for color and Adam for the rest

https://github.com/nerfstudio-project/gsplat/blob/main/examples/simple_trainer.py#L217-L224

I see, so it will automatically loop through all optimizers and check if any of the param group names match the paramdict?

in nerfstudio there is a separate torch.nn.Optimizer for each param group, so we should check if passing in a list of those works

List of optimizers also works as long as the param group has a "name" field.

The logic is that it will loop over the list of optimizers and every parameter group in it to find a match, based on the "name" field

Actually, maybe its more user friendly to just use a dict optimizer and force each parameter has a seperate optimizer?

kerrj · 2024-07-11T06:51:46Z

gsplat/strategy.py

+        # Refine GSs every this steps
+        refine_every: int = 100,
+        # Use absolute gradients for GS splitting
+        absgrad: bool = False,


if this is true maybe we want to automatically set the grow_grad2d to .008 otherwise the number of gaussians will explode? or maybe warn the user?

Not sure if warning is a good idea. on different scenes this threshold could be very different. Auto set is also not ideal as we want to give the user the freedom to adjust that.

I think the best we can do is probably note it clearly in the docstring.

Or, another idea is to create another class AbsGradStrategy that inherit from this class and set the default value to 0.008.

kerrj · 2024-07-11T06:59:41Z

gsplat/strategy.py

+    def _refine_duplicate(self, mask: Tensor):
+        device = mask.device
+        sel = torch.where(mask)[0]
+        for optimizer in self._optimizers:


do we want some sort of helper function for this? looping through optimizers and adding/pruning/removing is a common pattern that all strategies will need

maybe base level class functions that add/remove/duplicate an arbitrary set of indices for gaussians?

so these lines could then call self.duplicate_in_optimizer(index_mask) or something

I tried once but didn't end up creating anything elegant. Open for suggestions!

These functions in splatfacto try to do something like this, so any subclass could also call dup_in_optim and have the same effect

I think the body of the function would need to be different but the general interface of taking in a per-gaussian mask could work for removing/duplicating, or specifying how many to append

The nerfstudio way of implementation is mixing the duplication, splitting, and removing operations together so that the remove on optimizer can be done once. But the downside is that it's hard to read and modify. I would like to keep each operation isolated. But yeah maybe there is a way to share some code between them. I just dont have a concrete implementation idea.

put reusable functions into gsplat/strategy/ops.py!

gsplat/relocation.py

liruilong940607 · 2024-07-12T10:42:13Z

Preview Doc is online: https://669107b6a724ef1f2c186e67--inquisitive-speculoos-27b42b.netlify.app/apis/strategy

docs/source/apis/strategy.rst

rtabrizi and others added 18 commits July 7, 2024 17:04

[DRAFT] gsplat refactor (#269)

da01272

* initial commit, rough design of modules * moved strategy.py into gsplat dir * latest commit for 1:1 meeting * removed ops, losses, and strategies now inherit from an abstract 'Strategy' class * GSs init handled in private method

format

d4b6cf0

class Gaussian3D

3d68ebc

WIP: _strategy

2132d5d

simple_trainer no longer uses Gaussians object and instead uses Param…

e4fa44e

…eterDict with newly-added 'activations' dict

simple_trainer no longer uses Gaussians object and instead uses Param…

e10dc7c

…eterDict with newly-added 'activations' dict

created mcmc and revised strategy, simplified training scripts

0b9f760

update _strategy

c952363

Merge branch 'rtabrizi/refactor' of https://github.com/nerfstudio-pro…

1e60b9d

…ject/gsplat into rtabrizi/refactor

minor

29438a0

implemented ruilongs' suggested changes, no longer using separate cla…

4187c15

…ss for revised ADS heuristics

added docstrings, removed overriding activations attribute, fixed dev…

5c76a1c

…ice argument passing

mcmc and default are all trainable

b178fda

cleanup

2f536f2

minor

2da3eb9

minor

675e98b

minor

7617542

minor

2692398

liruilong940607 changed the title ~~[Refactor] Abstract out the densification logics (Original and MCMC) from the train script into gsplat library~~ Abstracting out the GS heuristics into a Strategy class Jul 11, 2024

Ruilong Li added 3 commits July 11, 2024 08:05

hardcode activations into strategy

4d962b2

minor cleanup

295b3f9

minor

5851782

liruilong940607 requested review from vye16, kerrj, brentyi, rtabrizi and maturk July 11, 2024 06:25

kerrj reviewed Jul 11, 2024

View reviewed changes

Ruilong Li added 16 commits July 12, 2024 01:37

stateless strategy

82e14b0

minor cleanup

8817495

revised_opacity

e2ad45f

_update_param_with_optimizer

9d7b456

individual ops

9d2ccf3

means3d -> means

cc8eef7

starting docs

b722f72

docs

85a58cd

docs

c1af65e

docs

860f73a

docs

d3d7b15

docs

9b43ea6

minor fix to pass test

67f439b

move BINOMS to strategy state

bcc843c

nits

ec94f3b

nit fix

024d72c

liruilong940607 commented Jul 12, 2024

View reviewed changes

gsplat/relocation.py Show resolved Hide resolved

liruilong940607 requested review from kerrj and brentyi July 12, 2024 10:41

viewcode

c9adf06

rtabrizi reviewed Jul 15, 2024

View reviewed changes

docs/source/apis/strategy.rst Outdated Show resolved Hide resolved

rtabrizi reviewed Jul 15, 2024

View reviewed changes

docs/source/apis/strategy.rst Outdated Show resolved Hide resolved

Ruilong Li added 5 commits July 17, 2024 22:06

relocation func cleanup

9ed34dc

nits in docs

825acf4

cleanup imports

568bae6

minor cleanup

9bf0a25

support screen size pruning / splitting

232b513

liruilong940607 merged commit 68ff2c1 into main Jul 17, 2024
2 checks passed

liruilong940607 deleted the rtabrizi/refactor branch July 17, 2024 23:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Abstracting out the GS heuristics into a Strategy class #278

Abstracting out the GS heuristics into a Strategy class #278

liruilong940607 commented Jul 11, 2024 •

edited

Loading

kerrj Jul 11, 2024

liruilong940607 Jul 11, 2024

kerrj Jul 11, 2024

kerrj Jul 11, 2024

liruilong940607 Jul 11, 2024

liruilong940607 Jul 11, 2024

kerrj Jul 11, 2024

liruilong940607 Jul 11, 2024

liruilong940607 Jul 11, 2024

kerrj Jul 11, 2024

kerrj Jul 11, 2024

kerrj Jul 11, 2024

liruilong940607 Jul 11, 2024

kerrj Jul 11, 2024

kerrj Jul 11, 2024

liruilong940607 Jul 11, 2024 •

edited

Loading

liruilong940607 Jul 12, 2024

liruilong940607 commented Jul 12, 2024

Abstracting out the GS heuristics into a Strategy class #278

Abstracting out the GS heuristics into a Strategy class #278

Conversation

liruilong940607 commented Jul 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liruilong940607 Jul 11, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liruilong940607 commented Jul 12, 2024

liruilong940607 commented Jul 11, 2024 •

edited

Loading

liruilong940607 Jul 11, 2024 •

edited

Loading