Add np.random.choice fallback for many Gaussians exceeding torch.multinomial limits

[soskek]

The current MCMCStrategy (relocate or sample_add) raises an error due to torch.multinomial when the number of Gaussians exceeds 2^24 = 16,777,216 (e.g., torch version is 2.0.1+cu117).

... gsplat/strategy/mcmc.py", line 158, in _relocate_gs
    relocate(
  File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File ".../.local/lib/python3.10/site-packages/gsplat/strategy/ops.py", line 228, in relocate
    sampled_idxs = torch.multinomial(probs, n, replacement=True)
RuntimeError: number of categories cannot exceed 2^24

This Pull Request introduces a fallback mechanism in the _multinomial_sample function to handle cases where the number of elements exceeds the limit. This is achieved by switching to np.random.choice when necessary, ensuring that the function can handle large element sets without encountering runtime errors.
Although np.random.choice may be slightly slower than torch.multinomial, it is not expected to become a bottleneck as long as refine_every is sufficiently large (e.g., 100; default).

[soskek]

It passes this simple test using pytest.

import math
import torch
from torch import nn, optim
import torch.testing

from gsplat.strategy.ops import sample_add, relocate


def _make_binoms():
    n_max = 51
    binoms = torch.zeros((n_max, n_max))
    for n in range(n_max):
        for k in range(n + 1):
            binoms[n, k] = math.comb(n, k)
    return binoms


def _test_sample_add(N):
    params = nn.ParameterDict({
        "opacities": nn.Parameter(torch.randn(N).cuda()),
        "scales": nn.Parameter(torch.randn(N, 3).cuda())
    })
    optimizers = {
        "opacities": optim.SGD([params["opacities"]], lr=0.01),
        "scales": optim.SGD([params["scales"]], lr=0.01)
    }
    state = {}
    binoms = _make_binoms().cuda()

    # Call the function
    sample_add(params, optimizers, state, n=5, binoms=binoms)

    # Assert the size of the params has increased correctly
    assert len(params["opacities"]) == N + 5
    assert len(params["scales"]) == N + 5


def _test_relocate(N):
    params = nn.ParameterDict({
        "opacities": nn.Parameter(torch.randn(N).cuda()),
        "scales": nn.Parameter(torch.randn(N, 3).cuda())
    })
    optimizers = {
        "opacities": optim.SGD([params["opacities"]], lr=0.01),
        "scales": optim.SGD([params["scales"]], lr=0.01)
    }
    state = {}
    mask = torch.zeros(N, dtype=torch.bool).cuda()
    mask[:5] = 1  # Mark the first 5 as "dead"
    binoms = _make_binoms().cuda()

    # Call the function
    relocate(params, optimizers, state, mask, binoms)

    # Assert that the sizes remain consistent and relocation has occurred
    assert len(params["opacities"]) == N
    assert len(params["scales"]) == N


def test_sample_add():
    _test_sample_add(N=30)
    _test_sample_add(N=2**24)
    _test_sample_add(N=2**24 + 1)


def test_relocate():
    _test_relocate(N=30)
    _test_relocate(N=2**24)
    _test_relocate(N=2**24 + 1)

[liruilong940607]

LGTM! Thanks for looking into this!

[liruilong940607]

Could you please format the code with black?

pip install black==22.3.0
black . gsplat/ tests/ examples/ profiling/

[soskek]

Thank you!