DGL工具系列(一)：用DGL实现pageRank算法

柔情只为你懂 2023-07-13 03:11 21阅读 0赞

import networkx as nx
    import matplotlib.pyplot as plt
    import torch
    import dgl
    
    N=6#100个节点
    DAMP=0.85#阻尼系数
    K=10#迭代次数
    
    g=dgl.DGLGraph()
    g.add_nodes(6)
    
    #添加边的关系
    g.add_edges([1,1,1,2,2,4,0,5,5],[2,3,4,3,5,3,1,3,4])#1,2 1,3 1,4 2,3 2,5 4,3 ,0,1 5,3 5,4 
    
    nx.draw(g.to_networkx(),node_size=220,node_color=[[.88, .5, .156,]],with_labels=True)
    plt.show()
    
    #定义pagerank算法最重要的两个变量  pagerank值和边权重
    g.ndata["pagerank"]=torch.ones(N)/N#节点的pagerank值
    g.ndata["degree"]=g.out_degrees(g.nodes()).float()#初始分配  每个节点的所有出边权重一样
    
    #定义消息函数  节点的pagerank值/节点的出度 所得到的值  即为要传递给目标节点的值
    def pagerankMessageFunc(edges):
        return {"pv" :edges.src["pagerank"]/edges.src["degree"]}#边源节点的pr/deg
    
    #d定义reduce函数  对邻居节点信息进行聚合
    def pagerankReduceFunc(nodes):
        msgs=torch.sum(nodes.mailbox["pv"],dim=1)
        pv=(1-DAMP)/N+msgs*DAMP#pageRank的总公式  更新当前节点的pagerank值
        return {"pv":pv}
    
    #注册以上两个函数
    g.register_message_func(pagerankMessageFunc)
    g.register_reduce_func(pagerankReduceFunc)
    
    #实现pagerank迭代
    def  pagerankIteration(g):
        for u,v in zip(*g.edges()):
            g.send((u,v))
        for v in g.nodes():
            g.recv(v)
    
    
    pagerankIteration(g)
    print(g.ndata["pagerank"])

第二种方式：

import dgl.function as fn
    
    def pagerank_builtin(g):
        g.ndata['pagerank'] = g.ndata['pagerank'] / g.ndata['degree']
        g.update_all(message_func=fn.copy_src(src='pagerank', out='m'),
                     reduce_func=fn.sum(msg='m',out='m_sum'))
        g.ndata['pagerank'] = (1 - DAMP) / N + DAMP * g.ndata['m_sum']
    
    for i in range(5):
        pagerank_builtin(g)
        print(g.ndata["pagerank"])