Group Normalization/ optimization of current layernom and instance norm

[dhorka]

I saw in the documentation that pytorch_geometric has the implementation of batch normalization,layer normalization, instance normalization but the group normalization is missing. Is there any plans to have the group normalization implemented in pytorch geometric?

[rusty1s]

The GroupNorm is indeed missing. Are you interested in contributing this feature?

[dhorka]

I haven't tried yet, but I was thinking to create the sparse tensor using the batch just to be able to use the sparse tensor. But I have to think more detail about it an check the current available operations with this sparse tensor.

[rusty1s]

Instead of using SparseTensor, you may be able to swap out the scatter calls with segment_csr calls from torch-scatter.

[dhorka]

I will check it this weekend, segment_csr should be faster than scatter operations?

[rusty1s]

Yes, it should be faster.

[dhorka]

I will check that during this week and I will come back to you

[dhorka]

I have tried as you suggested to use segment_csr. This is the implementation that I have done:

class LayerNorm(torch.nn.Module):
    def __init__(self, in_channels, eps=1e-5, affine=True):
        super(LayerNorm, self).__init__()

        self.in_channels = in_channels
        self.eps = eps

        if affine:
            self.weight = Parameter(torch.Tensor([in_channels]))
            self.bias = Parameter(torch.Tensor([in_channels]))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        ones(self.weight)
        zeros(self.bias)

    def forward(self, x: Tensor, batch: OptTensor = None) -> Tensor:
        if batch is None:
            x = x - x.mean()
            out = x / (x.std(unbiased=False) + self.eps)

        else:
            batch_size = int(batch.max()) + 1

            norm = degree(batch, batch_size, dtype=x.dtype).clamp_(min=1)
            norm = norm.mul_(x.size(-1)).view(-1, 1)
            
            # sorting indx
            indx = torch.argsort(batch, dim=0, descending=False)
            batch = batch[indx]
            x = x[indx]
            
            # creating ptr
            u , count = torch.unique(batch, return_counts=True)
            batch_ptr = [torch.tensor(0, device=batch.device)]
            for i, c in enumerate(count):
                batch_ptr.append(batch_ptr[i] + c)
            batch_ptr = torch.stack(batch_ptr).to(batch.device)
            
            
            mean = segment_csr(x, batch_ptr,
                           reduce='add').sum(dim=-1, keepdim=True) / norm

            x = x - mean[batch]
            
            var = segment_csr(x * x, batch_ptr,
                            reduce='add').sum(dim=-1, keepdim=True)
            
            var = var / norm

            out = x / (var + self.eps).sqrt()[batch]

        if self.weight is not None and self.bias is not None:
            out = out * self.weight + self.bias

        return out

    def __repr__(self):
        return f'{self.__class__.__name__}({self.in_channels})'

I have checked the speed up on my use case, and I didn't notice to much speed up. I think that maybe the problem is the time that I am saving with the segment_csr is spended during the creation of the ptr. Do you think that it could be done more efficient this part?

Thanks!

[rusty1s]

Ah, I see. Sorry for misunderstanding you. The problem is that you are operating on edge-level tensors, while ptr is node-level.

As a workaround, you can compute ptr from batch quite fast via:

torch.bucketize(torch.arange(batch_size, device=batch.device, batch)

[dhorka]

I assume that in this scenario there is not way to avoid the sort that I am doing at the beggining, isn't? I will check about your proposal I didn't know about this torch method.

[rusty1s]

I don't think the sort is necessary, as batch is sorted by default.

[dhorka]

Thanks! I will check it and I will come back with some feedback. For the moment what I had to do it is modify a little the torch.bucketize that you provided, just to obtain the same results that I was obtaining with my implementation:

batch_ptr = torch.cat([torch.tensor([0], device=batch.device),                                                                                    
                                    torch.bucketize(torch.arange(batch_size, device=batch.device), batch, right=True)])

[dhorka]

I have check it and there is an speedup but its still very slow in comparision to the batchnorm. This weekend I will try to prepare an example in colab to show a similar scenario to the one that I am using it. Also, I will try to think another way to do that, lets see if we can speedup more these kinds of normalizations.