Update forward & backward for rendered alpha image

[Zhuoyang-Pan]

I added the forward and backward for rendered alpha image. This could be helpful if we want to train on datasets with masks/RGBA images.

Specifically, for each call of RasterizeGaussians.apply or NDRasterizeGaussians.apply now we would have two outputs out_img (H * W * 3) and out_alpha(H * W). For $\alpha$ gradient for each gaussian, I calculate it as

$$ \begin{aligned} \frac{\partial{\mathcal{L}}}{ \partial{\alpha_n}} &= \sum_{k=1}^{K} \frac{\partial{\mathcal{L}}}{ \partial{C_i(k)}} \frac{\partial{C_i(k)}}{\partial{\alpha_n}} + \frac{\partial{\mathcal{L}}}{ \partial{A_i}} \frac{\partial{A_i}}{\partial{\alpha_n}} \\ &= \sum_{k=1}^{K} \frac{\partial{\mathcal{L}}}{ \partial{C_i(k)}} (c_n(k) \cdot T_n - \frac{\sum_{m > n} c_m \alpha_m T_m}{1-\alpha_n}) + \frac{\partial{\mathcal{L}}}{ \partial{A_i}} (T_n - \frac{\sum_{m > n} \alpha_m T_m}{1-\alpha_n}) \\ &= \sum_{k=1}^{K} \frac{\partial{\mathcal{L}}}{ \partial{C_i(k)}} (c_n(k) \cdot T_n - \frac{\sum_{m > n} c_m \alpha_m T_m}{1-\alpha_n}) + \frac{\partial{\mathcal{L}}}{ \partial{A_i}} \frac{T_{final}}{1-\alpha_n} \\ \end{aligned} $$

where $C_i$, $A_i$ are output colors, alpha respectively, $K$ is the number of channels of colors. The documents are also updated for reference of function calls.

I've also tested on following image where the left part has alpha=0.5 and the right part has alpha=1.

[kerrj]

I think ruilong removed the cuda() for these, could you take them out too?

[Zhuoyang-Pan]

fixed but should .cuda()s still be in NDRasterize backwards?

[kerrj]

Sorry for the late review! this looks great, can we add an optional flag to the rasterize forward function called return_alpha that defaults to False for backward compatibility though? After that change it looks good to merge

[maturk]

@Zhuoyang-Pan any ideas how this could be used with an input dataset that does contain masks to automatically segment the masked regions out?

[Zhuoyang-Pan]

@Zhuoyang-Pan any ideas how this could be used with an input dataset that does contain masks to automatically segment the masked regions out?

I am thinking of adding a L1 mask loss if the input dataset contains mask, but let me finish this PR first :)

[kerrj]

looks good, thanks!