Implementation of Message Passing of PointGNN

[andreasbinder]

Hi guys!
I am trying to implement the message passing that is presented in the paper "Point-GNN: Graph Neural Network for 3D Object Detection in a Point Cloud" .
The scheme looks like this :

The respective dataset it is working on is a point cloud (KITTI).
This is my first (conceptual) implementation; I know it is missing the correct channel dimension etc, it is more about where to put what :)

class PointGNNConv(MessagePassing):

  def __init__(
      self,
      in_channels: Union[int, Tuple[int, int]],
      out_channels: int,
      aggr: Optional[Union[str, List[str], Aggregation]] = "max",
  ):
      self.in_channels = in_channels
      self.out_channels = out_channels


      super().__init__(aggr)

      # calculate off-set delta x
      self.lin_h = nn.Linear(self.in_channels, self.out_channels)
      
      # calculate edge update
      self.lin_f = nn.Linear(self.in_channels, self.out_channels)

      # calculate final state update
      self.lin_g = nn.Linear(self.in_channels, self.out_channels)


  def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj, 
              size: Size = None) -> Tensor:


      out = self.propagate(edge_index, x=x, size=size)
      
      out = self.lin_g(out) + x

      return out


  def message(self, x_j: Tensor, x_i: Tensor, x0_j, x0_i, edge_index) -> Tensor:
      s_i = x_i 
      s_j = x_j 

      delta_x = self.lin_h(s_i)

      e = self.lin_f(torch.cat([x0_j - x0_i +delta_x, s_j]))
      return e

From my understanding, I receive e from the message function. I receive the states s_i and s_j from x through the propagate function. However, I do not only want to have the states for alle message pairs but also the (initial) coordinates x0_j and x0_i and I struggle how to implement this in the cleanest way. I thought about passing propagate a tensor that contains for every state also its respective coordinate.
Thank you!

[rusty1s]

I think this looks pretty good, but you are missing the initial coordinates in your function (we usually call them pos in PyG). So that would now look something like:

def forward(self, x, pos, edge_index):
    return x + self.lin(self.propagate(edge_index, x=x, pos=pos))

def message(self, pos_j, pos_i, x_i, x_j):
    return torch.cat([pos_j - pos_i + self.lin(x_i), x_j], dim=-1)

[andreasbinder]

Amazing, thank you so much!
Maybe two little follow-up questions:

1. Where can I find more information about those (implicit?) parameters like pos?
2. Once it is all tested, would the PointGNN layer be of interest for you by lets say making a pull request?

[rusty1s]

1. pos is not a parameter but an input argument of the model, and should be available on any mesh/point cloud dataset in PyG. See here for an example.
2. Definitely :)