我的任务是实现一个基于图神经网络的模型来模拟广度优先搜索算法。我正在进行某种迭代节点预测,因此我将节点以排序方式(增加节点 ID)放入 BFS 队列中,以便网络可以学习。(与随机未排序队列相比)
我正在使用 Deep Graph Library 和 python 来完成这项任务。如何在神经网络中实现 BFS 算法?我对此并不陌生,不确定将 BFS 函数放在我的神经网络中的哪个位置。
这是我当前的功能代码
'''Assign features to nodes or edges'''
def assign_features(num_nodes):
G.ndata['feat'] = torch.eye(num_nodes)
'''Function to build random graph'''
def build_rand_nx_graph(num_nodes, num_edges):
G = dgl.rand_graph(num_nodes,num_edges)
nx_G = G.to_networkx().to_undirected()
return nx_G
# for visualization purpose
'''Function to draw graph'''
def draw_graph(nx_G):
pos = nx.kamada_kawai_layout(nx_G)
nx.draw(nx_G, pos, with_labels=True, node_color=[[.7, .7, .7]])
'''Function to create BFS algorithm (deterministic part of the assignment)'''
def bfs(graph, root):
visited, queue = set(), collections.deque([root])
visited.add(root)
while queue:
# Dequeue a vertex from queue
vertex = queue.popleft()
print(str(vertex) + " ", end="")
# If not visited, mark it as visited, and
# enqueue it
neigh = graph[vertex]
neighbour = neigh.sort()
for neighbour in graph[vertex]:
if neighbour not in visited:
visited.add(neighbour)
queue.append(neighbour)
#main function
if __name__ == '__main__':
num_nodes = 50
num_edges = 60
nx_graph = build_rand_nx_graph(num_nodes,num_edges)
graph = nx.to_dict_of_lists(nx_graph)
#graph = nx.to_dict_of_lists(nx_G)
#graph = {0: [1, 2, 5, 6, 4], 1: [2], 2: [3], 3: [1, 2], 4:[0], 5:[0], 6:[0]}
print("Following is Breadth First Traversal: ")
bfs(graph, 0)
draw_graph(nx_graph)