解决PyG 报错 from torch_geometric.nn.pool.topk_pool import topk, filter_adj

zoe_ya

已于 2024-04-12 22:47:02 修改

阅读量2.3k

点赞数 23

CC 4.0 BY-SA版权

文章标签：图神经网络

于 2023-11-26 10:14:01 首次发布

本文链接：https://blog.youkuaiyun.com/Constantiny/article/details/134624540

文章讲述了在使用PyTorch的PyG库构建图神经网络时遇到关于topk_pool模块导入错误的问题，原因在于版本更新导致的语法变化。作者提供了通过替换和自定义函数解决此问题的方法，以及相关源码分析。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题：

使用Pytorch 的 PyG 搭建图神经网络报错

can not import topk, filter_adj from torch_geometric.nn.pool.topk_pool

解决

版本问题语法变化
topk => SelectTopk
filter_adj => FilterEdges

from torch_geometric.nn.pool.connect import FilterEdges
from torch_geometric.nn.pool.select import SelectTopK

发现替换后不可以
于是进去看SelectTopK\FilterEdges 源码
发现里面有 topk, filter_adj 方法但是直接 import 也不能用
于是手动写函数出来再 layers.py 里即可运行

def topk(
        x: Tensor,
        ratio: Optional[Union[float, int]],
        batch: Tensor,
        min_score: Optional[float] = None,
        tol: float = 1e-7,
) -> Tensor:
    if min_score is not None:
        # Make sure that we do not drop all nodes in a graph.
        scores_max = scatter(x, batch, reduce='max')[batch] - tol
        scores_min = scores_max.clamp(max=min_score)

        perm = (x > scores_min).nonzero().view(-1)
        return perm

    if ratio is not None:
        num_nodes = scatter(batch.new_ones(x.size(0)), batch, reduce='sum')

        if ratio >= 1:
            k = num_nodes.new_full((num_nodes.size(0),), int(ratio))
        else:
            k = (float(ratio) * num_nodes.to(x.dtype)).ceil().to(torch.long)

        x, x_perm = torch.sort(x.view(-1), descending=True)
        batch = batch[x_perm]
        batch, batch_perm = torch.sort(batch, descending=False, stable=True)

        arange = torch.arange(x.size(0), dtype=torch.long, device=x.device)
        ptr = cumsum(num_nodes)
        batched_arange = arange - ptr[batch]
        mask = batched_arange < k[batch]

        return x_perm[batch_perm[mask]]


def filter_adj(
        edge_index: Tensor,
        edge_attr: Optional[Tensor],
        node_index: Tensor,
        cluster_index: Optional[Tensor] = None,
        num_nodes: Optional[int] = None,
) -> Tuple[Tensor, Optional[Tensor]]:
    num_nodes = maybe_num_nodes(edge_index, num_nodes)

    if cluster_index is None:
        cluster_index = torch.arange(node_index.size(0),
                                     device=node_index.device)

    mask = node_index.new_full((num_nodes,), -1)
    mask[node_index] = cluster_index

    row, col = edge_index[0], edge_index[1]
    row, col = mask[row], mask[col]
    mask = (row >= 0) & (col >= 0)
    row, col = row[mask], col[mask]

    if edge_attr is not None:
        edge_attr = edge_attr[mask]

    return torch.stack([row, col], dim=0), edge_attr

参考官方文档

https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/pool/topk_pool.html