topk函数的详细解释

本文主要是torch_geometric库中的topk函数的代码进行解释，该函数的在库中的路径为 torch_geometric.nn.pool,topk_pool.topk。在torch_geometric的文档中并没有有关于该函数的任何记录，可能是因为该函数属于比较底层的操作。但是在某些情况下，我们可能不得不去了解底层函数的执行过程来更好地帮助我解决实际问题，或者构造属于自己的module。
函数输入：x：节点的得分，ratio：保留节点的比例，batch：图分配向量
函数输出：perm：保留的节点编码
函数功能：保留每张子图中得分最高的ratio%的节点

'''对topk函数的解析'''
def topk(x, ratio, batch, min_score=None, tol=1e-7):
    if min_score is not None:
        # Make sure that we do not drop all nodes in a graph.
        scores_max = scatter_max(x, batch)[0][batch] - tol
        scores_min = scores_max.clamp(max=min_score)

        perm = (x > scores_min).nonzero(as_tuple=False).view(-1)
    else:
        # 向量中每一个元素表示了子图中节点的个数
        num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0)
        batch_size, max_num_nodes = num_nodes.size(0), num_nodes.max().item()
        # 节点累积向量，第k个元素表示了前k-1个子图的总节点数
        cum_num_nodes = torch.cat(
            [num_nodes.new_zeros(1),
             num_nodes.cumsum(dim=0)[:-1]], dim=0)

        index = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
        # 统一每一张子图中的节点数，并将其节点编码映射到相应的编码空间中
        index = (index - cum_num_nodes[batch]) + (batch * max_num_nodes)

        dense_x = x.new_full((batch_size * max_num_nodes, ),
                             torch.finfo(x.dtype).min)
        dense_x[index] = x
        dense_x = dense_x.view(batch_size, max_num_nodes)
        # 对每张子图中节点的分数，按照降序排序
        _, perm = dense_x.sort(dim=-1, descending=True)
        """
        将每张子图中的节点编码映射回最开始的编码空间中，在向量中可能存在相同的编码。
        这是因为两个向量出现在不同的子图中,不影响选取每张子图得分最高的前k个节点。
        """
        perm = perm + cum_num_nodes.view(-1, 1)
        perm = perm.view(-1)
        # 选取每张子图的前k个节点
        k = (ratio * num_nodes.to(torch.float)).ceil().to(torch.long)
        mask = [
            torch.arange(k[i], dtype=torch.long, device=x.device) +
            i * max_num_nodes for i in range(batch_size)
        ]
        mask = torch.cat(mask, dim=0)
        # 掩码操作，perm保存了最终保留的节点的编码
        perm = perm[mask]

    return perm