Questions about VNMaxPool

[guochengqian]

Dear @FlyingGiraffe,

Thanks for your impressive work.

I have a question about the VNMaxPool layer. You used the argmax to select the idx of the pooled features. However, the idx is not differentiable, and thus the self.map_to_dir layer will not be updated anyway.

vnn/models/vn_layers.py

Line 139 in f3197e2

d = self.map_to_dir(x.transpose(1,-1)).transpose(1,-1)

Could you explain this a little bit? Maybe I am wrong.

Thank you very much!

[Julien-Gustin]

Hi @guochengqian,
I have the same problem, have you found how to solve this problem?

Thank you. (cc. @FlyingGiraffe )

[guochengqian]

@Julien-Gustin Hi thanks for ping. No I quit on this problem.

[chrockey]

Hi, @FlyingGiraffe. Thanks for open-sourcing your codes.
May I ask you to answer this issue?

[udaykamal20]

I tried using softargmax but still ended up with the same error (in DDP mode); therefore finally went for the softmaxpool instead of hardmaxpool. This ensures differentiability and update to the self.map_to_dir.

class VNMaxPool(nn.Module):

    def __init__(self, in_channels):
        super(VNMaxPool, self).__init__()
        self.map_to_dir = nn.Linear(in_channels, in_channels, bias=False)

    def forward(self, x):
        '''
        x: point features of shape [B, N_feat, 3, N_samples, ...]
        '''
        d = self.map_to_dir(x.transpose(1,-1)).transpose(1,-1)
        dotprod = (x*d).sum(2, keepdims=True)
        sm = torch.nn.functional.softmax(dotprod*100., dim=-1) #100 or any large value to scale the input so that the max value dominates over others
        x_max = (x*sm).sum(-1)
        return x_max