Implementation of PointTransformer

[QuanticDisaster]

Hi,

While looking for recent architectures to try for 3D semantic segmentation, I came across PointTransformer, so I gave it a try to implement it with the Message Passing scheme.

The implementation is based on PointTransformer from Hengshuang Zhao and all (2020) https://arxiv.org/abs/2012.09164, more specifically in the 3.1, 3.2 and 3.3 sections and features an attention layer.

import torch

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops
from torch.nn import Sequential, ReLU, Linear
from torch.nn.functional import softmax
from torch_cluster import knn_graph

from typing import Union
from torch_geometric.typing import OptTensor, PairOptTensor, PairTensor, Adj

from torch import Tensor


class TransformerLayer(MessagePassing):
    def __init__(self, dim_feat, dim_out, pos_mlp_hidden_dim=64, attn_mlp_hidden_dim=2 ):
        super(TransformerLayer, self).__init__(aggr='add')  # "Add" aggregation
        
        self.delta = Sequential(Linear(3,pos_mlp_hidden_dim), ReLU(), Linear(pos_mlp_hidden_dim,dim_out))
        self.gamma = Sequential(Linear(dim_out, dim_out* attn_mlp_hidden_dim), ReLU(), Linear(dim_out * attn_mlp_hidden_dim,dim_out))
        self.alpha = Linear(dim_feat, dim_out)
        self.phi =   Linear(dim_feat, dim_out)
        self.psi =   Linear(dim_feat, dim_out)
        

    def forward(self, x: Union[OptTensor, PairOptTensor],
                pos: Union[Tensor, PairTensor], edge_index: Adj) -> Tensor:
        """"""
        
        if not isinstance(x, tuple):
            x: PairOptTensor = (x, x)

        if isinstance(pos, Tensor):
            pos: PairTensor = (pos, pos)

        return self.propagate(edge_index, x=x, pos=pos)


    def message(self, x_i,x_j, pos_i, pos_j):
        
        #query is x_i, key is x_j, value is the substraction of the two
        out_delta = self.delta((pos_j - pos_i) ) #compute position encoding
        
        out_phi_psi = self.phi(x_i) - self.psi(x_j) # compute point wise transformation
        out_gamma = self.gamma(out_phi_psi + out_delta) #compute importance factor
                
        msg = softmax(out_gamma,1) * ( self.alpha(x_j) + out_delta) #normalize the result with softmax and multiply by features transformed by alpha
        
        return msg

if __name__ == '__main__':
    pos = torch.randn(20,3)
    x = torch.randn(20,2)
    
        
    #PointTransformer test
    edge_index = knn_graph(pos, 1, loop = True)
    layer = TransformerLayer(dim_feat=2, dim_out=2)
    result = layer(x,pos,edge_index)

Unfortunately, there is no official implementation, so some details are not known such as the intermediate dimensions for the PointTransformer layer or when features cardinality is augmented or reduced.

This is why I also had a look at these github repositories /~https://github.com/POSTECH-CVLab/point-transformer and /~https://github.com/qq456cvb/Point-Transformers as well as /~https://github.com/lucidrains/point-transformer-pytorch (PointTransformer layer only)

The transformer block itself :

class TransformerBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels, hidden_channels):
        super(TransformerBlock, self).__init__()
        self.linear_in = Sequential( Linear(in_channels + 3, hidden_channels), ReLU())
        self.linear_out = Sequential( Linear(out_channels + 3, out_channels), ReLU())
        self.transformer = TransformerLayer(dim_feat = hidden_channels, dim_out = out_channels)
        
    def forward(self, x, pos, edge_index):
        x = self.linear_in(torch.cat([pos,x], dim=1))
        x = self.transformer(x, pos, edge_index)
        x = self.linear_out(torch.cat([pos,x], dim=1))
        
        return x

However, the network does not seem to support points only input, at least one feature is mandatory (or I didn't implement it correctly, which could be the case feel free to correct it !). While I tried to implement the full architecture, I didn't try it on the S3DIS or shapenet dataset.

I would be glad to hear your thoughts on it and see if someones wants to do a full example !

[rusty1s]

Thanks. This is great. Are you interested in contributing this layer directly? We can work on an example for it in addition or later.

Note that we can usually simply pass in dummy node features in case none exists, e.g., a simply 1 for every node via the ToConstant(1) transform. This enforces the model to solely learn from relative spatial relations, instead of injecting absolute coordinates via x = self.linear_in(torch.cat([pos,x], dim=1)).

[QuanticDisaster]

Thanks. This is great. Are you interested in contributing this layer directly? We can work on an example for it in addition or later.

Do you mean make a pull request ? I can, however I am new to contributing in general

Note that we can usually simply pass in dummy node features in case none exists, e.g., a simply 1 for every node via the ToConstant(1) transform. This enforces the model to solely learn from relative spatial relations, instead of injecting absolute coordinates via x = self.linear_in(torch.cat([pos,x], dim=1)).

Indeed. I will have a try with this, but the last time I tried tro bypass it this way with GCN if i remember well, results were terrible

[rusty1s]

Yes, for contributing, you can simply

1. fork PyG and clone it locally
2. Create a new branch: git checkout -b point_transformer
3. Make your modifications, e.g., by adding your code to torch_geometric/nn/conv/point_transformer_conv.py
4. Push your changes
5. Go to GitHub and create a PR

GCN does indeed not well with constant features, as GCNConv has no way to discriminate between neighbors. This does not hold for layers operating on point cloud data, as they can still weight the importance of neighbors based on their relative spatial positions.

[QuanticDisaster]

Alright I gave it a try and created the PR request

GCN does indeed not well with constant features, as GCNConv has no way to discriminate between neighbors. This does not hold for layers operating on point cloud data, as they can still weight the importance of neighbors based on their relative spatial positions.

I see, I didn't think about it, thx for pointing out !