Allow for >1 batch size in Splatfacto

[akristoffersen]

WIP, preliminary testing makes it look like it's working but I would want to make sure.

[akristoffersen]

this will fail with masks / depth images

[kerrj]

Hey Alex! this is super cool; especially in MCMC which doesn't require gradient thresholds at all. #3216 might have mildly broken parts of this PR since it merged in parallel dataloading, but it shouldn't be too bad; let us know if you want any help fixing conflicts!

[hardikdava]

@akristoffersen I think you might want to modify step_post_backward to support for batch processing. So densification and pruning can be intact as it is.

[akristoffersen]

Works with masks now

As expected, I noticed a almost 2x increase in rays/s with a batch size of two, and a very slight performance drop with a batch size of 1 compared to baseline (50.1 M rays/sec -> 48 M rays/sec)

[akristoffersen]

@hardikdava do you mean that the tuning might be different for the thresholds? Yeah, I don't know exactly what to do there. maybe someone else has an opinion?

Some quick stats on the poster dataset.
Orange: Baseline (batch size of 1)
Blue: BS = 5
Red: BS = 10

so the splitting / densification outcomes are affected by batch size.

Similarly, train rays/sec do start higher due to the larger batch size, but go down as you'd expect with the higher number of gaussians.

Some good news, with a higher batch I do see the training loss hitting better values quicker as the batch size increases.

[hardikdava]

@akristoffersen currently, densification, splitting and culling are implemented inside strategy and logic is based on step. Since batch training can skip some of the functions. I am not sure but you might want to modify the strategy function by calling the batch size e.g. loop for the batch size. So that all densification logic will be performed for all steps, otherwise some of the functions will be skipped.

In simple words, suppose the batch size is 2, opacity reset needs to be applied at every 3000th step. So it should happen at every 1500th steps according to batch size. But according to your current implementation it will be applied at every 3000th steps but actually it will be 6000th step (batch size * step).