JUST代码理解

JUST代码理解

这是我学习异构图嵌入所看的第一篇文章，将代码理解写在博客为了日后之便利，疏漏之处，各位谅解

import argparse
import numpy as np
import networkx as nx
import random
import math
from gensim.models import Word2Vec

argparse时python自带的命令行参数解析包。其它模块在用到处解释。

详细见argparse简要用法总结

def parse_args():
	parser = argparse.ArgumentParser(description="Just")

	parser.add_argument('--input')

	parser.add_argument('--node_types')

	parser.add_argument('--output')

	parser.add_argument('--dimensions' , type=int)

	parser.add_argument('--walk_length' , type=int)

	parser.add_argument('--num_walks' , type=int)

	parser.add_argument('--window-size' , type=int)
	
	parser.add_argument('--alpha' , type=float)

	parser.add_argument('--workers' , default=1)

	return parser.parse_args()

使用参数解析器从命令行获取训练模型所需要的参数，包括：

alpha停留在同一域概率，dimensions训练得到的结点向量维度, input='Datasets/dblp.edgelist'异构图, node_types异构边的类型，num_walks在每个结点游走次数, output目标输出文件，walk_length随机游走长度, window_size选区上下文窗口大小。

在命令行输入命令后，函数进入python程序的模拟入口：

if __name__ == "__main__":
	args = parse_args()
	main(args)

得到参数，传递并运行main()

def main(args):
	G = nx.read_edgelist(args.input)
	heterg_dictionary = generate_node_types()
	walks = generate_walks(G, args.num_walks, args.walk_length,heterg_dictionary)
	print 'Starting training .. '
	model = Word2Vec(walks, size=args.dimensions, window=args.window_size,
                     min_count=0, workers=args.workers)
	print 'Finished training .. '
	model.save_word2vec_format(args.output)

2行读取异构图，3行读取异构图的异构结构，在3行后加入print语句：

{'a': ['p'], 'p': ['a', 't', 'v'], 't': ['p'], 'v': ['p']}

可以看到得到的时异构边的集合。

第4行生成随机游走

#此函数从图中使用JUST方法，为每个节点生成num_walks个长度
#为walk_length的随机游走
def generate_walks(G, num_walks, walk_length, heterg_dictionary):
  print 'Generating walks .. '
  walks = []
  nodes = list(G.nodes())

  for cnt in range(num_walks):
    random.shuffle(nodes)
    for node in nodes:
		just_walks = dblp_generation(G, walk_length, heterg_dictionary,
                                     start=node)
		walks.append(just_walks)
  print 'Walks done .. '
  return walks

JUST随机游走函数：

#  number of memorized domains = 1
def dblp_generation(G, path_length, heterg_dictionary, start=None):
    
    path = []
	#将传入节点设为游走初始节点
    path.append(start)
    
    cnt = 1
    homog_length = 1
    no_next_types = 0
    
    heterg_probability = 0
    #随机游走长度
    while len(path) < path_length:
		if no_next_types == 1:
			break
		
		cur = path[-1]
		homog_type = []
		heterg_type = []
		#获得当前节点的同构边和异构边集合
		for node_type in heterg_dictionary:
			#cur[0]字符串首位，节点前缀即节点种类
			if cur[0] == node_type:
				homog_type = node_type
				heterg_type = heterg_dictionary[node_type]
		#计算JUMP的概率，为1-Pr_{stay}，停留概率Pr_{stay}等于alpha的(游走中当前节点往前的长度)次方
		heterg_probability = 1 - math.pow(args.alpha, homog_length)
		r = random.uniform(0, 1)
		#生成随机数
		next_type_options=[]
		#跳转
		if r <= heterg_probability:
			#demo样例中使用的m=1,即跳转的下一个节点域当前节点种类不同即可
			for heterg_type_iterator in heterg_type:
				next_type_options.extend([e for e in G[cur] if (e[0] == heterg_type_iterator)])
			#当前节点没有异构边，从同构边中选取
			if not next_type_options:
				next_type_options = [e for e in G[cur] if (e[0] == homog_type)]
		else:
			next_type_options = [e for e in G[cur] if (e[0] == homog_type)]
			if not next_type_options:
				for heterg_type_iterator in heterg_type:
					next_type_options.extend([e for e in G[cur] if (e[0] == heterg_type_iterator)])	
		if not next_type_options:
			no_next_types = 1
			break
			
		next_node = random.choice(next_type_options)
		path.append(next_node)
		if next_node[0] == cur[0]:
			homog_length = homog_length + 1
		else:
			 homog_length = 1		

    return path

第6行将的到的随机游走采样投喂给word2vec，得到节点的向量表示，训练结束。

JUST的确是一篇简单易懂好入门的优秀论文。