torch_geometric.nn.conv.PointGNNConv

class PointGNNConv(mlp_h: Module, mlp_f: Module, mlp_g: Module, **kwargs)[source]

Bases: MessagePassing

The PointGNN operator from the “Point-GNN: Graph Neural Network for 3D Object Detection in a Point Cloud” paper.

\[ \begin{align}\begin{aligned}\Delta \textrm{pos}_i &= h_{\mathbf{\Theta}}(\mathbf{x}_i)\\\mathbf{e}_{j,i} &= f_{\mathbf{\Theta}}(\textrm{pos}_j - \textrm{pos}_i + \Delta \textrm{pos}_i, \mathbf{x}_j)\\\mathbf{x}^{\prime}_i &= g_{\mathbf{\Theta}}(\max_{j \in \mathcal{N}(i)} \mathbf{e}_{j,i}) + \mathbf{x}_i\end{aligned}\end{align} \]

The relative position is used in the message passing step to introduce global translation invariance. To also counter shifts in the local neighborhood of the center node, the authors propose to utilize an alignment offset. The graph should be statically constructed using radius-based cutoff.

Parameters:
  • mlp_h (torch.nn.Module) – A neural network \(h_{\mathbf{\Theta}}\) that maps node features of size \(F_{in}\) to three-dimensional coordination offsets \(\Delta \textrm{pos}_i\).

  • mlp_f (torch.nn.Module) – A neural network \(f_{\mathbf{\Theta}}\) that computes \(\mathbf{e}_{j,i}\) from the features of neighbors of size \(F_{in}\) and the three-dimensional vector \(\textrm{pos_j} - \textrm{pos_i} + \Delta \textrm{pos}_i\).

  • mlp_g (torch.nn.Module) – A neural network \(g_{\mathbf{\Theta}}\) that maps the aggregated edge features back to \(F_{in}\).

  • **kwargs (optional) – Additional arguments of torch_geometric.nn.conv.MessagePassing.

Shapes:
  • input: node features \((|\mathcal{V}|, F_{in})\), positions \((|\mathcal{V}|, 3)\), edge indices \((2, |\mathcal{E}|)\),

  • output: node features \((|\mathcal{V}|, F_{in})\)

forward(x: Tensor, pos: Tensor, edge_index: Union[Tensor, SparseTensor]) Tensor[source]

Runs the forward pass of the module.

reset_parameters()[source]

Resets all learnable parameters of the module.