torch_geometric.nn.pool.radius

radius(x: Tensor, y: Tensor, r: float, batch_x: Optional[Tensor] = None, batch_y: Optional[Tensor] = None, max_num_neighbors: int = 32, num_workers: int = 1, batch_size: Optional[int] = None) → Tensor[source]

为y中的每个元素找到x中距离r以内的所有点。

import torch
from torch_geometric.nn import radius

x = torch.tensor([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.tensor([[-1.0, 0.0], [1.0, 0.0]])
batch_y = torch.tensor([0, 0])
assign_index = radius(x, y, 1.5, batch_x, batch_y)
Parameters:

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

  • y (torch.Tensor) – 节点特征矩阵 \(\mathbf{Y} \in \mathbb{R}^{M \times F}\).

  • r (float) – The radius.

  • batch_x (torch.Tensor, optional) – 批次向量 \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), 它将每个 节点分配给一个特定的示例。(默认: None)

  • batch_y (torch.Tensor, optional) – 批次向量 \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M\), 它将每个 节点分配给特定的示例。(默认: None)

  • max_num_neighbors (int, optional) – The maximum number of neighbors to return for each element in y. (default: 32)

  • 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)

  • batch_size (int, optional) – The number of examples \(B\). Automatically calculated if not given. (default: None)

Return type:

torch.Tensor

警告

radius() 的 CPU 实现在使用 max_num_neighbors 时偏向某些象限。 考虑将 max_num_neighbors 设置为 None 或在继续之前将输入移动到 GPU。