fix: affine gradient propagation

[pierremerriaux-leddartech]

hi,
In affine transformation module, gradients don't propagate to mlp layers and images embeddings, they are always equal to zero. The affine transformation is only learnt in last layer biases.
it is due to mlp weight and biases zeros initialization. I propose to initialize only the last layer to zero, to enable gradient propagation. Kept the last layer to 0, able to start train close to identity matrix

thanks
Pierre

[georghess]

@pierremerriaux-leddartech did you find this to help? Currently trying to evaluate OmniRe on Pandaset, but I still get clear artifacts from different camera settings even with this fix, see sky regions in images below.

Another issue with the current approach is that test images use the average embedding over all cameras, not just the embeddings for the corresponding camera. However, changing the averaging didn't improve results notably.

@ziyc Any input on this? I want to make a fair comparison between OmniRe and other methods I'm benchmarking.

[pierremerriaux-leddartech]

Hi @georghess,
Nice to see again ! With my PR, gradient will be propagated properly, but yes we have also some artifact of exposure in sky from different camera. bilateral grid helped us more about that.

affine:

bilateral:

We have to finish affine correction tuning, but not sure that exposure is constant enough over image area and it is for that bilateral grid help.
thanks

[georghess]

@pierremerriaux-leddartech thanks for the quick reply! I think one of the issues is also that the sky node (EnvLight) has multiple order of magnitudes higher lr than the affine transform, leading to the model quickly changing the sky color and never recovering from it. But haven't had time to tune this to verify the hypothesis

[pierremerriaux-leddartech]

thanks @georghess , yes probably, we will test that during tuning phase.
thanks

[pierremerriaux-leddartech]

Hi @szhang963 ,
Sure, for now we use the bilateral grid nerfstudio implementation in splatfacto.
thanks