dgl源码阅读笔记（3）——DeepWalk

图神经网络开源库dgl阅读笔记

文章目录

dgl源码阅读笔记（3）——DeepWalk
- 图神经网络开源库dgl阅读笔记
- @[TOC](文章目录)
前言
一、DeepWalk简单回顾
二、浏览代码类
- 1.class DeepwalkTrainer
- 2.class SkipGramModel.__init__()
- - def init_emb2pos_index
  - def init_emb2neg_index
- class SkipGramModel的其他函数
- - def fast_learn
- def adam
- 3. class DeepwalkSampler
- - def random_walk
三、代码执行流程
- 1. 主函数启动
- - fast_train

前言

博客内容为个人阅读dgl代码笔记，仅供个人使用，还有许多不足需要发现改正。

一、DeepWalk简单回顾

通过对每一个采样点进行随机游走，生成多条具有相同长度的游走序列，通过Skipgram模型对游走序列进行训练学习，可以让每个节点的embedding得到梯度下降的更新。在这其中还会用到负采样技术。

二、浏览代码类

1.class DeepwalkTrainer

初始化函数如下：
–norm 是否在embedding后进行归一化
–dim embedding的维度默认128
–window_size 窗口大小默认5
–num_walks 每个节点随机走的次数默认是10
–negative 是否负采样
–batch_size 默认是128
–walk_length 默认80
–neg_weight 负采样的权重
–lr 学习率
–true_neg False的时候代表不要求负采样精度很高，使用faster_sample。如果是True，就要保证负采样里没有FN数据，而生成TN的数据进行负采样。
–fast_neg 和上面的true_neg 是互斥的，True的时候代表精度可以不用很高，采用fast_neg方法

    def __init__(self, args):""" Initializing the trainer with the input arguments """self.args = argsself.dataset = DeepwalkDataset(net_file=args.data_file,map_file=args.map_file,walk_length=args.walk_length,window_size=args.window_size,num_walks=args.num_walks,batch_size=args.batch_size,negative=args.negative,gpus=args.gpus,fast_neg=args.fast_neg,ogbl_name=args.ogbl_name,load_from_ogbl=args.load_from_ogbl,)self.emb_size = self.dataset.G.number_of_nodes()# 初始化one-hot向量self.emb_model = None

还有以下运行需要用到的函数：

判断在什么设备商运行，并且重要的是在这个函数里初始化了SkipGramModel：

def init_device_emb(self):

决定启动fast_train_mp或者fast_train：

def train(self):

启动多设备下的训练：

def fast_train_mp(self):

fast_train_mp的子进程

def fast_train_sp(self):

单设备下的训练：（本次运行涉及到）

def fast_train(self):

2.class SkipGramModel.init()

这里的参数每一个都先存到self的同名变量中

    """ Negative sampling based skip-gram """def __init__(self, emb_size, emb_dimension,walk_length,window_size,batch_size,only_cpu,only_gpu,mix,neg_weight,negative,lr,lap_norm,fast_neg,record_loss,norm,use_context_weight,async_update,num_threads,):

接下来首先是对u和v的embedding进行随机初始化
并把embedding限制在initrange(1/128)的正态分布内
再用0填充weight中的空值

        # initialize the device as cpuself.device = torch.device("cpu")# content embeddingself.u_embeddings = nn.Embedding(self.emb_size, self.emb_dimension, sparse=True)# context embeddingself.v_embeddings = nn.Embedding(self.emb_size, self.emb_dimension, sparse=True)# initialze embeddinginitrange = 1.0 / self.emb_dimensioninit.uniform_(self.u_embeddings.weight.data, -initrange, initrange)  # 把embedding限制在initrange(1/128)的正态分布内init.constant_(self.v_embeddings.weight.data, 0) # 用0填充weight中的空值，虽然我怀疑是否用得上

初始化lookup_table，打散成-6到6区间，用于生成logsigmoid_table ，并把首尾赋值为0和1(sigmoid两端)

        # lookup_table is used for fast sigmoid computingself.lookup_table = torch.sigmoid(torch.arange(-6.01, 6.01, 0.01))self.lookup_table[0] = 0.self.lookup_table[-1] = 1.if self.record_loss:self.logsigmoid_table = torch.log(torch.sigmoid(torch.arange(-6.01, 6.01, 0.01)))self.loss = []

下一步是对正负样本的节点序列选取。

        # indexes to select positive/negative node pairs from batch_walksself.index_emb_posu, self.index_emb_posv = init_emb2pos_index(self.walk_length,self.window_size,self.batch_size)self.index_emb_negu, self.index_emb_negv = init_emb2neg_index(self.walk_length,self.window_size,self.negative,self.batch_size)

正负节点样本选取调用的两个函数，我以正样本为例，进入了下面的函数

def init_emb2pos_index

该函数返回出正样本，即存在的边对应的u，v向量
按照定义来说，在同一个batch下的窗口内的两个节点就是正采样的节点，所以按照此三重循环，可以取出所有满足的节点对。b和j用于控制采样是在同一个batch下，满足条件的i和j对应的就是u,v两个节点。

def init_emb2pos_index(walk_length, window_size, batch_size):idx_list_u = []idx_list_v = []for b in range(batch_size):for i in range(walk_length):for j in range(i-window_size, i):if j >= 0:idx_list_u.append(j + b * walk_length)idx_list_v.append(i + b * walk_length)for j in range(i + 1, i + 1 + window_size):if j < walk_length:idx_list_u.append(j + b * walk_length)idx_list_v.append(i + b * walk_length)# [num_pos * batch_size]index_emb_posu = torch.LongTensor(idx_list_u)index_emb_posv = torch.LongTensor(idx_list_v)return index_emb_posu, index_emb_posv

def init_emb2neg_index

对于负采样进行比较
同样是输出u,v，三重循环中，只是选取了同样数量的u，乘上negtive的倍数。
因为是fast_sample，所以对v的选取不是很准确，而是生成这样大小的列表list(range(batch_size * walk_length)) * negative * window_size * 2 ，这样肯定会导致超出所需。
下一步idx_list_v = idx_list_v[:len(idx_list_u)] 把v列表控制在u列表的大小中。

def init_emb2neg_index(walk_length, window_size, negative, batch_size):idx_list_u = []for b in range(batch_size):for i in range(walk_length):for j in range(i-window_size, i):if j >= 0:idx_list_u += [i + b * walk_length] * negativefor j in range(i+1, i+1+window_size):if j < walk_length:idx_list_u += [i + b * walk_length] * negativeidx_list_v = list(range(batch_size * walk_length))\* negative * window_size * 2random.shuffle(idx_list_v)idx_list_v = idx_list_v[:len(idx_list_u)]# [bs * walk_length * negative]index_emb_negu = torch.LongTensor(idx_list_u)index_emb_negv = torch.LongTensor(idx_list_v)return index_emb_negu, index_emb_negv

class SkipGramModel的其他函数

然后回到SkipGram.__init__的下一步，决定是否初始化上下文权重矩阵
初始化优化器、梯度

        if self.use_context_weight:self.context_weight = init_weight(self.walk_length,self.window_size,self.batch_size)# adamself.state_sum_u = torch.zeros(self.emb_size)self.state_sum_v = torch.zeros(self.emb_size)# gradients of nodes in batch_walksself.grad_u, self.grad_v = init_empty_grad(self.emb_dimension,self.walk_length,self.batch_size)

还需要注意到以下函数：

设置异步更新子进程：

    def create_async_update(self):

终止异步线程更新：

def finish_async_update(self):

多线程共享参数

    def share_memory(self):

决定设备参数，并把lookup_table，grad等移到上面

    def set_device(self, gpu_id):

把参数移到gpu上

    def all_to_device(self, gpu_id):

用预处理的lookup_table，返回sigmoid和fast_logsigmoid值

    def fast_sigmoid(self, score):idx = torch.floor((score + 6.01) / 0.01).long()return self.lookup_table[idx]def fast_logsigmoid(self, score):idx = torch.floor((score + 6.01) / 0.01).long()return self.logsigmoid_table[idx]

前向传播：

    def forward(self, pos_u, pos_v, neg_v):

存下训练后的节点embedding，为了以后的使用：

    def save_embedding(self, dataset, file_name):def save_embedding_pt(self, dataset, file_name):def save_embedding_pt_dgl_graph(self, dataset, file_name):def save_embedding_txt(self, dataset, file_name):

进行随机游走序列的提取：

    def fast_learn(self, batch_walks, neg_nodes=None):

接下来展开讲解这个比较重要的fastlearn函数：

def fast_learn

以下是函数说明。
不用前向传播自动计算梯度，用lookup_table计算实现sigmoid
并用u_embedding[i] += (label - score) * v_embedding[j]的方式更新特征向量
其中的batch_walks就是传入的walks，也就是128*80（batch有128个点，路径是80步）的矩阵，下图也是对该函数的一些说明，比较好理解

    def fast_learn(self, batch_walks, neg_nodes=None):"""Learn a batch of random walks in a fast way. It has the following features:1. It calculating the gradients directly without the forward operation.2. It does sigmoid by a looking up table.Specifically, for each positive/negative node pair (i,j), the updating procedure is as following:score = self.fast_sigmoid(u_embedding[i].dot(v_embedding[j]))# label = 1 for positive samples; label = 0 for negative samples.u_embedding[i] += (label - score) * v_embedding[j]v_embedding[i] += (label - score) * u_embedding[j]Parameters----------batch_walks list : a list of node sequenceslr float : current learning rateneg_nodes torch.LongTensor : a long tensor of sampled true negative nodes. If neg_nodes is None,then do negative sampling randomly from the nodes in batch_walks as an alternative.Usage example-------------batch_walks = [torch.LongTensor([1,2,3,4]), torch.LongTensor([2,3,4,2])])lr = 0.01neg_nodes = None"""

首先要明确的是，batch_walks(128,80)是batch_num和walklength，所以是所有出现在当前batch中的所有walk的节点集合
接下来规定了学习率，存储当前序列到nodes中
Self.u_embeddings和self.v_embeddings是利用torch.nn.embedding随机生成的矩阵，目前还没有进行训练，特别的是用的不是one-hot向量，而是随机生成的，且维度为(1138499, 128)整个数据集的维度乘特征向量维度
emb_u,emb_v对应的是将128*80的nodes放入，且对应节点编号寻找对应的128长度的特征向量，这样每一行对应一个节点的初始向量。

        lr = self.lr# [batch_size, walk_length]if isinstance(batch_walks, list):nodes = torch.stack(batch_walks)elif isinstance(batch_walks, torch.LongTensor):nodes = batch_walksif self.only_gpu:nodes = nodes.to(self.device)if neg_nodes is not None:neg_nodes = neg_nodes.to(self.device)emb_u = self.u_embeddings(nodes).view(-1, self.emb_dimension).to(self.device)emb_v = self.v_embeddings(nodes).view(-1, self.emb_dimension).to(self.device)

接下来开始进行正采样样本的训练和更新过程：
首先判断batch是否是最后一个batch，然后将正采样的节点对(u,v)序列存到下面的数据结构中

Emb_pos_u就是从emb_u中，以第0维度为索引，查找所有在index_emb_posu中的节点，这时就有98560的维度
98560 = 128 * (800 - 2 * (5 + 5 * 4 / 2))

        ## Postivebs = len(batch_walks)if bs < self.batch_size:index_emb_posu, index_emb_posv = init_emb2pos_index(self.walk_length, self.window_size, bs)index_emb_posu = index_emb_posu.to(self.device)index_emb_posv = index_emb_posv.to(self.device)else:index_emb_posu = self.index_emb_posuindex_emb_posv = self.index_emb_posv# num_pos: the number of positive node pairs generated by a single walk sequence# [batch_size * num_pos, dim]emb_pos_u = torch.index_select(emb_u, 0, index_emb_posu)emb_pos_v = torch.index_select(emb_v, 0, index_emb_posv)

接下来开始计算每一对节点特征向量相乘，因为是正采样，需要达到1为最好，score被定义为如下的模式

        pos_score = torch.sum(torch.mul(emb_pos_u, emb_pos_v), dim=1) # torch.mul是对两个结构一样的矩阵，每一个对应位置上的数字相乘，所以不改变维度，然后再按列求和，这样列维度就变成了1pos_score = torch.clamp(pos_score, max=6, min=-6)# 最大值和最小值被限制在了[-6,6]维度中，不是归一化，而是超过范围的值都赋值为边界值，因为超过这个范围进入sigmoid的值相差不会有多少，防止过平滑# [batch_size * num_pos, 1]score = (1 - self.fast_sigmoid(pos_score)).unsqueeze(1) # unsqueeze用于增加维度，本来是行向量的score,变成了列向量if self.record_loss:self.loss.append(torch.mean(self.fast_logsigmoid(pos_score)).item())

接下来计算正采样梯度，可能会考虑到拉普拉斯归一化

        # [batch_size * num_pos, dim]if self.lap_norm > 0:grad_u_pos = score * emb_pos_v + self.lap_norm * (emb_pos_v - emb_pos_u)grad_v_pos = score * emb_pos_u + self.lap_norm * (emb_pos_u - emb_pos_v)else:grad_u_pos = score * emb_pos_vgrad_v_pos = score * emb_pos_uif self.use_context_weight:if bs < self.batch_size:context_weight = init_weight(self.walk_length,self.window_size,bs).to(self.device)else:context_weight = self.context_weightgrad_u_pos *= context_weightgrad_v_pos *= context_weight

然后初始化self.grad，并对grad_u和grad_v进行更新
利用到了这个函数:index_add_
对于x.index_add_(0, index, source) 0是对应的维度(行），x是要更新的矩阵，index对应每次更新x的哪一行，source从上往下遍历的矩阵，每次从最上面取一行加到Index对应的x行中。

        # [batch_size * walk_length, dim]if bs < self.batch_size:grad_u, grad_v = init_empty_grad(self.emb_dimension, self.walk_length, bs)grad_u = grad_u.to(self.device)grad_v = grad_v.to(self.device)else:self.grad_u = self.grad_u.to(self.device)self.grad_u.zero_()self.grad_v = self.grad_v.to(self.device)self.grad_v.zero_()grad_u = self.grad_ugrad_v = self.grad_vgrad_u.index_add_(0, index_emb_posu, grad_u_pos)grad_v.index_add_(0, index_emb_posv, grad_v_pos)

下面是负采样更新梯度，大同小异

        ## Negativeif bs < self.batch_size:index_emb_negu, index_emb_negv = init_emb2neg_index(self.walk_length, self.window_size, self.negative, bs)index_emb_negu = index_emb_negu.to(self.device)index_emb_negv = index_emb_negv.to(self.device)else:index_emb_negu = self.index_emb_neguindex_emb_negv = self.index_emb_negvemb_neg_u = torch.index_select(emb_u, 0, index_emb_negu)if neg_nodes is None:emb_neg_v = torch.index_select(emb_v, 0, index_emb_negv)else:emb_neg_v = self.v_embeddings.weight[neg_nodes].to(self.device)# [batch_size * walk_length * negative, dim]neg_score = torch.sum(torch.mul(emb_neg_u, emb_neg_v), dim=1)neg_score = torch.clamp(neg_score, max=6, min=-6)# [batch_size * walk_length * negative, 1]score = - self.fast_sigmoid(neg_score).unsqueeze(1)if self.record_loss:self.loss.append(self.negative * self.neg_weight * torch.mean(self.fast_logsigmoid(-neg_score)).item())grad_u_neg = self.neg_weight * score * emb_neg_vgrad_v_neg = self.neg_weight * score * emb_neg_ugrad_u.index_add_(0, index_emb_negu, grad_u_neg)if neg_nodes is None:grad_v.index_add_(0, index_emb_negv, grad_v_neg)

接下来是更新阶段：
在这里将刚才计算出的梯度放到adam优化器中更新grad

        ## Updatenodes = nodes.view(-1) # -1代表的是没有考虑到的剩下的维度数量，比如这里就是将128*80的矩阵变成了10240的向量# use adam optimizergrad_u = adam(grad_u, self.state_sum_u, nodes, lr, self.device, self.only_gpu)grad_v = adam(grad_v, self.state_sum_v, nodes, lr, self.device, self.only_gpu)if neg_nodes is not None:grad_v_neg = adam(grad_v_neg, self.state_sum_v, neg_nodes, lr, self.device, self.only_gpu)

最后利用grad_u和grad_v将u和v的embedding进行更新，

        if not self.async_update:self.u_embeddings.weight.data.index_add_(0, nodes.view(-1), grad_u)self.v_embeddings.weight.data.index_add_(0, nodes.view(-1), grad_v)            if neg_nodes is not None:self.v_embeddings.weight.data.index_add_(0, neg_nodes.view(-1), grad_v_neg)return

def adam

在计算梯度后，引入了adam优化器，处理后的grad才会给embedding更新用
关于Adam优化器原理可以参考这篇博客：简单认识Adam优化器

下面代码中，std_values 是vt，grad就是mt，比正规的Adam实现简单了一些

def adam(grad, state_sum, nodes, lr, device, only_gpu):""" calculate gradients according to adam """grad_sum = (grad * grad).mean(1)if not only_gpu:grad_sum = grad_sum.cpu()state_sum.index_add_(0, nodes, grad_sum) # cpustd = state_sum[nodes].to(device)  # gpustd_values = std.sqrt_().add_(1e-10).unsqueeze(1)grad = (lr * grad / std_values) # gpu

3. class DeepwalkSampler

整个类的代码如下：
其实是调用dgl.sampling.random_walk产生随机游走结果

class DeepwalkSampler(object):def __init__(self, G, seeds, walk_length):""" random walk sampler Parameter---------G dgl.Graph : the input graphseeds torch.LongTensor : starting nodeswalk_length int : walk length"""self.G = Gself.seeds = seedsself.walk_length = walk_lengthdef sample(self, seeds):walks = dgl.sampling.random_walk(self.G, seeds, length=self.walk_length-1)[0]return walks

def random_walk

返回值的第一个部分就是根据nodes序列生成的shape(nodes)个随机游走序列

def random_walk(g, nodes, *, metapath=None, length=None, prob=None, restart_prob=None,return_eids=False):

传参给了g, nodes,length，其余为默认值
这里面最重要的是
traces, eids, types = _CAPI_DGLSamplingRandomWalk(gidx, nodes, metapath, p_nd)
traces就是随机游走后的序列

    n_etypes = len(g.canonical_etypes)n_ntypes = len(g.ntypes)if metapath is None:if n_etypes > 1 or n_ntypes > 1:raise DGLError("metapath not specified and the graph is not homogeneous.")if length is None:raise ValueError("Please specify either the metapath or the random walk length.")metapath = [0] * lengthelse:metapath = [g.get_etype_id(etype) for etype in metapath]gidx = g._graphnodes = F.to_dgl_nd(utils.prepare_tensor(g, nodes, 'nodes'))metapath = F.to_dgl_nd(utils.prepare_tensor(g, metapath, 'metapath'))# Load the probability tensor from the edge framesif prob is None:p_nd = [nd.array([], ctx=nodes.ctx) for _ in g.canonical_etypes]else:p_nd = []for etype in g.canonical_etypes:if prob in g.edges[etype].data:prob_nd = F.to_dgl_nd(g.edges[etype].data[prob])if prob_nd.ctx != nodes.ctx:raise ValueError('context of seed node array and edges[%s].data[%s] are different' %(etype, prob))else:prob_nd = nd.array([], ctx=nodes.ctx)p_nd.append(prob_nd)# Actual random walkif restart_prob is None:traces, eids, types = _CAPI_DGLSamplingRandomWalk(gidx, nodes, metapath, p_nd)elif F.is_tensor(restart_prob):restart_prob = F.to_dgl_nd(restart_prob)traces, eids, types = _CAPI_DGLSamplingRandomWalkWithStepwiseRestart(gidx, nodes, metapath, p_nd, restart_prob)else:traces, eids, types = _CAPI_DGLSamplingRandomWalkWithRestart(gidx, nodes, metapath, p_nd, restart_prob)traces = F.from_dgl_nd(traces)types = F.from_dgl_nd(types)eids = F.from_dgl_nd(eids)return (traces, eids, types) if return_eids else (traces, types)

三、代码执行流程

1. 主函数启动

也就是DeepwalkTrainer开始启动，传入默认参数，建立出DeepwalkTrainer
然后进入train函数

    trainer = DeepwalkTrainer(args)trainer.train()

    def train(self):""" train the embedding """if len(self.args.gpus) > 1:self.fast_train_mp()else:self.fast_train()

接着跳转到fast_train中

fast_train

下面的一部分是对相关的参数进行初始化，这里唯一需要注意的num_pos（正采样个数）是考虑到序列前后节点不足window个邻居对应的数量。

    def fast_train(self):""" fast train with dataloader with only gpu / only cpu"""# the number of postive node pairs of a node sequencenum_pos = 2 * self.args.walk_length * self.args.window_size \- self.args.window_size * (self.args.window_size + 1)  # 后面减去的内容就是在随机游走序列两端的不够window长度的节点所缺少的训练pos数num_pos = int(num_pos)self.init_device_emb() # 根据传入的arg参数对device进行赋值# 和定义SkipGramModel()if self.args.async_update:self.emb_model.share_memory()self.emb_model.create_async_update()if self.args.count_params:sum_up_params(self.emb_model)sampler = self.dataset.create_sampler(0)dataloader = DataLoader(dataset=sampler.seeds,batch_size=self.args.batch_size,collate_fn=sampler.sample,shuffle=False,drop_last=False,num_workers=self.args.num_sampler_threads,)num_batches = len(dataloader)print("num batchs: %d\n" % num_batches)

需要注意到在 sampler = self.dataset.create_sampler(0)，调用了这个函数，接着返回了DeepwalkSampler类。
所以可以在下一行的dataloader = DataLoader()实现产生随机游走序列迭代器的功能。

    def create_sampler(self, i):""" create random walk sampler """return DeepwalkSampler(self.G, self.seeds[i], self.walk_length)

下面的代码，常规的torch.no_grad()，在此之内的代码不收梯度计算的影响
对dataloader里的每个batch的编号和128*80的向量序列进行遍历
进入到emb_model.fast_learn(walks)中，在这里进行向量的更新过程。
其他的操作就是输入输出的控制，判断之前是否进行过fast负采样，没有的话就进行np.random.choice的负采样，还有控制输出的语句。
这个代码就结束了。

        start_all = time.time()start = time.time()with torch.no_grad():max_i = num_batchesfor i, walks in enumerate(dataloader):# 对dataloader里的每个batch的编号和128*80的向量序列进行遍历if self.args.fast_neg:self.emb_model.fast_learn(walks)else:# do negative samplingbs = len(walks)neg_nodes = torch.LongTensor(np.random.choice(self.dataset.neg_table,bs * num_pos * self.args.negative,replace=True))self.emb_model.fast_learn(walks, neg_nodes=neg_nodes)if i > 0 and i % self.args.print_interval == 0:if self.args.print_loss:print("Batch %d training time: %.2fs loss: %.4f" \% (i, time.time() - start, -sum(self.emb_model.loss) / self.args.print_interval))self.emb_model.loss = []else:print("Batch %d, training time: %.2fs" % (i, time.time() - start))start = time.time()if self.args.async_update:self.emb_model.finish_async_update()print("Training used time: %.2fs" % (time.time() - start_all))if self.args.save_in_txt:self.emb_model.save_embedding_txt(self.dataset, self.args.output_emb_file)elif self.args.save_in_pt:self.emb_model.save_embedding_pt(self.dataset, self.args.output_emb_file)else:self.emb_model.save_embedding(self.dataset, self.args.output_emb_file)