Abstracting out the GS heuristics into a Strategy class (#278)

* [DRAFT] gsplat refactor (#269)

* 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

* class Gaussian3D

* WIP: _strategy

* simple_trainer no longer uses Gaussians object and instead uses ParameterDict with newly-added 'activations' dict

* simple_trainer no longer uses Gaussians object and instead uses ParameterDict with newly-added 'activations' dict

* created mcmc and revised strategy, simplified training scripts

* update _strategy

* minor

* implemented ruilongs' suggested changes, no longer using separate class for revised ADS heuristics

* added docstrings, removed overriding activations attribute, fixed device argument passing

* mcmc and default are all trainable

* cleanup

* minor

* minor

* minor

* minor

* hardcode activations into strategy

* minor cleanup

* minor

* seperate the strategy in to individual files

* change optimizers from List to Dict

* switch to dataclass

* stateless strategy

* minor cleanup

* revised_opacity

* _update_param_with_optimizer

* individual ops

* means3d -> means

* starting docs

* docs

* docs

* docs

* docs

* docs

* minor fix to pass test

* move BINOMS to strategy state

* nits

* nit fix

* viewcode

* relocation func cleanup

* nits in docs

* cleanup imports

* minor cleanup

* support screen size pruning / splitting

---------

Co-authored-by: Ryan Tabrizi <[email protected]>
Co-authored-by: Ruilong Li <[email protected]>
Co-authored-by: rtabrizi <[email protected]>

Author: 3 people

docs/source/apis/strategy.rst

@@ -0,0 +1,76 @@
+Densification Strategy
+===================================
+
+.. currentmodule:: gsplat
+
+In `gsplat`, we abstract out the densification and pruning process of the Gaussian 
+training into a strategy. A strategy is a class that defines how the Gaussian parameters
+(along with their optimizers) should be updated (splitting, pruning, etc.) during the training. 
+
+An example of the training workflow using :class:`DefaultStrategy` is like:
+
+.. code-block:: python
+
+    from gsplat import DefaultStrategy, rasterization
+
+    # Define Gaussian parameters and optimizers
+    params: Dict[str, torch.nn.Parameter] | torch.nn.ParameterDict = ...
+    optimizers: Dict[str, torch.optim.Optimizer] = ...
+
+    # Initialize the strategy
+    strategy = DefaultStrategy()
+
+    # Check the sanity of the parameters and optimizers
+    strategy.check_sanity(params, optimizers)
+
+    # Initialize the strategy state
+    strategy_state = strategy.initialize_state()
+
+    # Training loop
+    for step in range(1000):
+        # Forward pass
+        render_image, render_alpha, info = rasterization(...)
+
+        # Pre-backward step
+        strategy.step_pre_backward(params, optimizers, strategy_state, step, info)
+
+        # Compute the loss
+        loss = ...
+
+        # Backward pass
+        loss.backward()
+
+        # Post-backward step
+        strategy.step_post_backward(params, optimizers, strategy_state, step, info)
+
+A strategy will inplacely update the Gaussian parameters as well as the optimizers,
+so it has a specific expectation on the format of the parameters and the optimizers.
+It is designed to work with the Guassians defined as either a Dict of
+`torch.nn.Parameter <https://pytorch.org/docs/stable/generated/torch.nn.parameter.Parameter.html>`_
+or a
+`torch.nn.ParameterDict <https://pytorch.org/docs/stable/generated/torch.nn.ParameterDict.html>`_
+with at least the following keys: {"means", "scales", "quats", "opacities"}. On top of these attributes,
+an arbitrary  number of extra attributes are supported. Besides the parameters, it also
+expects a Dict of `torch.optim.Optimizer <https://pytorch.org/docs/stable/optim.html>`_
+with the same keys as the parameters, and each optimizer should correspond to only
+one learnable parameter.
+
+For example, the following is a valid format for the parameters and the optimizers
+that can be used with our strategies:
+
+.. code-block:: python
+
+    N = 100
+    params = torch.nn.ParameterDict{
+        "means": Tensor(N, 3), "scales": Tensor(N), "quats": Tensor(N, 4), "opacities": Tensor(N),
+        "colors": Tensor(N, 25, 3), "features1": Tensor(N, 128), "features2": Tensor(N, 64),
+    }
+    optimizers = {k: torch.optim.Adam([p], lr=1e-3) for k, p in params.keys()}
+
+Below are the strategies that are currently implemented in `gsplat`:
+
+.. autoclass:: DefaultStrategy
+    :members:
+
+.. autoclass:: MCMCStrategy
+    :members:

docs/source/conf.py

@@ -29,6 +29,7 @@
     "sphinx.ext.intersphinx",
     "sphinxcontrib.bibtex",
     "sphinxcontrib.video",
+    "sphinx.ext.viewcode",
 ]
 
 intersphinx_mapping = {