Question about spherical initialzation and training

[ryunuri]

Hi, thanks for sharing your code.

I've been trying out several things and found something weird.
When using sphere initialization of the vanilla MLP, I expected the initial shape to be a sphere.
If you render the outputs of the initialized model by setting val_check_interval=1, the images (rgb, normal, depth) indeed resemble a sphere.

However, the marching cubes fail with the following error message

vmin, vmax = mesh_coarse['v_pos'].amin(dim=0), mesh_coarse['v_pos'].amax(dim=0)
IndexError: amin(): Expected reduction dim 0 to have non-zero size.

I guess this means that the aabb cube is empty.

When I looked into the code, I found that the VanillaMLP does not initialize the constants of the layers, which is different from the initialization of the paper "SAL: Sign Agnostic Learning of Shapes from Raw Data".

I think the make_linear function should be as follows

def make_linear(self, dim_in, dim_out, bias, is_first, is_last):
    layer = nn.Linear(dim_in, dim_out)
    if self.sphere_init:
        if is_last:
            torch.nn.init.constant_(layer.bias, -bias)
            torch.nn.init.normal_(layer.weight, mean=math.sqrt(math.pi) / math.sqrt(dim_in), std=0.0001)
        elif is_first:
            torch.nn.init.constant_(layer.bias, 0.0)
            torch.nn.init.constant_(layer.weight[:, 3:], 0.0)
            torch.nn.init.normal_(layer.weight[:, :3], 0.0, math.sqrt(2) / math.sqrt(dim_out))
        else:
            torch.nn.init.constant_(layer.bias, 0.0)
            torch.nn.init.normal_(layer.weight, 0.0, math.sqrt(2) / math.sqrt(dim_out))
    else:
        torch.nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')
    
    if self.weight_norm:
        layer = nn.utils.weight_norm(layer)
    return layer   

Also, from forward and forward_level methods in class VolumeSDF

if 'sdf_activation' in self.config:
            sdf = get_activation(self.config.sdf_activation)(sdf + float(self.config.sdf_bias))

The if statement is True even when you simply set sdf_activation to None in the config, since it's still in the config. I found that this leads the sdf values to be all positive at the start of training. I just removed the sdf_activation in the config.

After changing this part and setting the bias of the SDF to 0.6, the initial model output is as follows:

And the result of marching cubes is indeed a sphere.

However, I found that by changing the model like this results in very poor training results.

After 1000 iterations,

Also, the mesh is completely broken

So, I guess you had a reason for this design choice? Otherwise, I think this might be the reason why training the model on my custom dataset fails.

[bennyguo]

Hi, thanks for the valuable comments!

About the initialization of bias in linear layers, I think the current implementation achieves the same effect as the original code you mentioned, as the bias is set to 0 in intermediate layers and set to -bias in the last layer which is done by sdf_bias in our implementation. The only difference is that I set sdf_bias=0 instead of sdf_bias=-0.5 in the original setting, which leads to a smaller sphere.

I also found that setting sdf_bias=-0.6 could result in bad training results as you showed. It seems like a training instability problem and could be solved by adopting learning rate warm-up like in the original NeuS implementation. I already updated the config file to support such a warm-up strategy, please have a try!

[bennyguo]

Here I compare sdf_bias=0, w/o warm-up and sdf_bias=-0.5, w/ warm-up on the Lego scene. It seems that the latter achieves higher quality:

sdf_bias=0, w/o warm-up

sdf_bias=-0.5, w/ warm-up

[ryunuri]

Thanks for your feedback!

I tried out the warm-up strategy and the quality was quite improved on my custom dataset.
However, I'm still having some difficulties on extracting a high quality mesh out of it.
I guess I'll try out some more experiments and share the results if I see some improvements.