torch_geometric.nn.unpool.knn_interpolate

knn_interpolate(x: Tensor, pos_x: Tensor, pos_y: Tensor, batch_x: Optional[Tensor] = None, batch_y: Optional[Tensor] = None, k: int = 3, num_workers: int = 1)

The k-NN interpolation from the “PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space” paper.

For each point \(y\) with position \(\mathbf{p}(y)\), its interpolated features \(\mathbf{f}(y)\) are given by

\[\mathbf{f}(y) = \frac{\sum_{i=1}^k w(x_i) \mathbf{f}(x_i)}{\sum_{i=1}^k w(x_i)} \textrm{, where } w(x_i) = \frac{1}{d(\mathbf{p}(y), \mathbf{p}(x_i))^2}\]

and \(\{ x_1, \ldots, x_k \}\) denoting the \(k\) nearest points to \(y\).

Parameters:

• x (torch.Tensor) – Node feature matrix \(\mathbf{X} \in \mathbb{R}^{N \times F}\).

• pos_x (torch.Tensor) – Node position matrix \(\in \mathbb{R}^{N \times d}\).

• pos_y (torch.Tensor) – Upsampled node position matrix \(\in \mathbb{R}^{M \times d}\).

• batch_x (torch.Tensor, optional) – Batch vector \(\mathbf{b_x} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node from \(\mathbf{X}\) to a specific example. (default: None)

• batch_y (torch.Tensor, optional) – Batch vector \(\mathbf{b_y} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node from \(\mathbf{Y}\) to a specific example. (default: None)

• k (int, optional) – Number of neighbors. (default: 3)

• num_workers (int, optional) – Number of workers to use for computation. Has no effect in case batch_x or batch_y is not None, or the input lies on the GPU. (default: 1)