torch_cluster.random_walk()函数

torch_cluster链接

作用：Samples random walks of length walk_length from all node indices in start in the graph given by (row, col).

1、 以官网给出的例子生成有向图

import matplotlib.pyplot as plt
import networkx as nx
row = [0, 1, 1, 1, 2, 2, 3, 3, 4, 4]
col = [1, 0, 2, 3, 1, 4, 1, 4, 2, 3]
g = nx.DiGraph()
g.add_edges_from(list(zip(row, col)))
nx.draw(g, with_labels=True) 
plt.show()

2、在图上进行随机游走

import torch
from torch_cluster import random_walk

row = torch.tensor([0, 1, 1, 1, 2, 2, 3, 3, 4, 4])
col = torch.tensor([1, 0, 2, 3, 1, 4, 1, 4, 2, 3])
start = torch.tensor([0, 1, 2, 3, 4])

walk = random_walk(row, col, start, walk_length=3)
print(walk)

返回列表：shape(节点个数, 长度+1)

tensor([[0, 1, 2, 4],
        [1, 3, 4, 2],
        [2, 4, 2, 1],
        [3, 4, 2, 4],
        [4, 3, 1, 0]])

每次运行结果都有随机性， 大家可以自己验证一下