How to calculate the mean of pos with respect to batch data.

[saedrna]

Hi,

I would like to ask how to get the mean with respect to the batch information. And then I can subtract it from the pos.
Functions like global_xxx_pool(x, batch) only output a single value for a batch, but I'd like to expand (or repeat) it.
Following is some modified code from the pointnet2_classification example.

Thanks,
Han

class GlobalSAModule(torch.nn.Module):
    def __init__(self, nn):
        super().__init__()
        self.nn = nn

    def forward(self, x, pos, batch):
        # pos -> (num_nodes, 3)
        pos_center = ... # get the center with respect to the batch
        pos_centered = pos - pos_center # (num_nodes, 3)
        
        x = self.nn(torch.cat([x, pos_centered], dim=1))
        # returns x as -> (B, D)
        x = global_max_pool(x, batch)
        pos = pos.new_zeros((x.size(0), 3))
        batch = torch.arange(x.size(0), device=batch.device)
        return x, pos, batch

[saedrna]

I answered my own question. I have found a function in pytorch index_select. I think this should be the solution.

B = 10
N = 100
x = torch.randn(N, 3)
b = torch.randint(0,B,(N,))
c = torch.randn(B, 3)
x = x - torch.index_select(c, 0, b)

[rusty1s]

Yes, this looks good! (sorry for the delay in answering)