torch_geometric.nn.pool.radius_graph

radius_graph(x: Tensor, r: float, batch: Optional[Tensor] = None, loop: bool = False, max_num_neighbors: int = 32, flow: str = 'source_to_target', num_workers: int = 1, batch_size: Optional[int] = None) → Tensor[source]

计算到给定距离内所有点的图边。

import torch
from torch_geometric.nn import radius_graph

x = torch.tensor([[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, 1.0]])
batch = torch.tensor([0, 0, 0, 0])
edge_index = radius_graph(x, r=1.5, batch=batch, loop=False)
Parameters:

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

  • r (float) – 半径。

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

  • loop (bool, optional) – If True, the graph will contain self-loops. (default: False)

  • max_num_neighbors (int, optional) – 每个元素在 y 中返回的最大邻居数。（默认值：32）

  • flow (str, optional) – The flow direction when using in combination with message passing ("source_to_target" or "target_to_source"). (default: "source_to_target")

  • num_workers (int, optional) – Number of workers to use for computation. Has no effect in case batch 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_graph() 的 CPU 实现在使用 max_num_neighbors 时对某些象限存在偏差。建议将 max_num_neighbors 设置为 None 或在继续之前将输入移动到 GPU。