桑运鑫作者上海交通大学硕士生学校图神经网络应用研究方向

大规模图训练调优指南

最近对一个大规模的图训练嵌入,发现相关的中文资料还是很欠缺的,把自己踩的一些坑记下来。本文主要针对 DGL [1] 和 PyTorch [2] 两个框架。

一、训练大规模图

对于大规模图不能像小图一样把整张图扔进去训练,需要对大图进行采样,即通过Neighborhood Sampling方法每次采样一部分输出节点,然后把更新它们所需的所有节点作为输入节点,通过这样的方式做 mini-batch 迭代训练。具体的用法可以参考官方文档中的 Chapter 6: Stochastic Training on Large Graphs [3]。

但是 GATNE-T [4] 中有一种更有趣的做法,即只把 DGL 作为一个辅助计算流的工具,提供Neighborhood SamplingMessage Passing等过程,把Node EmbeddingEdge Embedding等存储在图之外,做一个单独的Embedding矩阵。每次从 dgl 中获取节点的id之后再去Embedding矩阵中去取对应的 embedding 进行优化,以此可以更方便的做一些优化。

二、缩小图规模

从图的Message Passing过程可以看出,基本上所有的图神经网络的计算都只能传播连通图的信息,所以可以先用 connected_componets [5] 检查一下自己的图是否是连通图。如果分为多个连通子图的话,可以分别进行训练或者选择一个大小适中的 component 训练。

如果图还是很大的话,也可以对图整体做一次Neighborhood Sampling采样一个子图进行训练。

三、减小内存占用

对于大规模数据而言,如何在内存中存下它也是一件让人伤脑筋的事情。这时候采用什么样的数据结构存储就很关键了。首先是不要使用原生的list,使用np.ndarray或者torch.tensor。尤其注意不要显式的使用set存储大规模数据(可以使用set去重,但不要存储它)。

注意:四种数据结构消耗的内存之间的差别(比例关系)会随着数据规模变大而变大。

其次就是在PyTorch中,设置DataLoadernum_workers大于 0 时会出现内存泄露的问题,内存占用会随着 epoch 不断增大。查阅资料有两个解决方法:
  • 根据Num_workers in DataLoader will increase memory usage? [6],设置num_workers小于实际 cpu 数,亲测无效;

  • 根据 CPU memory gradually leaks when num_workers > 0 in the DataLoader [7],将原始list转为np.ndarray或者torch.tensor,可以解决。原因是:There is no way of storing arbitrary python objects (even simple lists) in shared memory in Python without triggering copy-on-write behaviour due to the addition of refcounts, everytime something reads from these objects.

四、减小显存消耗

对于大规模图嵌入而言,Embedding矩阵会非常大。在反向传播中如果对整个矩阵做优化的话很可能会爆显存。可以参考 pinsage [8] 的代码,设置Embedding矩阵的sparse = True,使用 SparseAdam [9] 进行优化。SparseAdam是一种为 sparse 向量设计的优化方法,它只优化矩阵中参与计算的元素,可以大大减少backward过程中的显存消耗。

此外,如果显存仍然不够的话,可以考虑将Embedding矩阵放到 CPU 上。使用两个优化器分别进行优化。

五、加快训练速度

对于大规模数据而言,训练同样要花很长时间。加快训练速度也很关键。加快训练速度方面主要在两个方面:加快数据预处理和提高 GPU 利用率。

在加快数据预处理中,大部分数据集样本之间都是独立的,可以并行处理,所以当数据规模很大的时候,一定要加大数据预处理的并行度。

不要使用 for 循环逐条处理,可以使用 multiprocess [10] 库开多进程并行处理。但是要注意适当设置processes,否则会出现错误OSError: [Errno 24] Too many open files

此外,数据处理好之后最好保存为pickle格式的文件,下次使用可以直接加载,不要再花时间处理一遍。

在提高 GPU 利用率上,如果 GPU 利用率比较低主要是两个原因:batch_size较小(表现为 GPU 利用率一直很低)和数据加载跟不上训练(表现为 GPU 利用率上升下降上升下降,如此循环)。解决方法也是两个:

  • 增大batch_size,一般来说 GPU 利用率和batch_size成正比;

  • 加快数据加载:设置DataLoaderpin_memory=True, 适当增大num_workers(注意不要盲目增大,设置到使用的 CPU 利用率到 90% 左右就可以了,不然反而可能会因为开线程的消耗拖慢训练速度)。

如果有多块 GPU 的话,可以参照 graphsage [11] 的代码进行多 GPU 训练。一般来说在单机多卡的情况下都可以得到线性加速。但是使用DistributedDataParallel有几个需要注意的问题:
  • 最好参照 graphsage [11] 中的代码而不是使用官方教程中的torch.multiprocessing.spawn函数开辟多进程。因为使用这个函数会不停的打印Using backend: pytorch,暂时还不清楚是什么原因。
  • DataParallel一样,DistributedDataParallel会对原模型进行封装,如果需要获取原模型model的一些属性或函数的话,需要将model替换为model.module
  • 在使用DistributedDataParallel时,需要根据 GPU 的数量对batch_sizelearning rate进行调整。根据 Should we split batch_size according to ngpu_per_node when DistributedDataparallel [12] ,简单来说就是保持batch_sizelearning rate的乘积不变,因为我们多 GPU 训练一般不改batch_size,所以使用了多少 GPU 就要把learning rate扩大为原来的几倍。
  • 如何使DistributedDataParallel支持Sparse Embedding? 可以参考我在 PyTorch 论坛上的回答 DistributedDataParallel Sparse Embeddings [13],设置torch.distributed.init_process_group中的backend=gloo即可,现在版本(1.6 以及 nightly)的 PyTorch 在nccl中仍然不支持Sparse Embedding。关于这个问题的最新进展可以看这个 PR:Sparse allreduce for ProcessGroupNCCL [14]

最后,PyTorch 1.6 中提供了混合精度训练 amp [15] 的 API,可以方便的调用,通过将一部分操作从torch.float32 (float)变为torch.float16 (half)来加快训练速度。但是它并不支持Sparse Embedding,如果你的模型中包含Sparse向量的话,最好不要使用。

参考文献

[1] https://www.dgl.ai/
[2] https://pytorch.org/
[3] https://docs.dgl.ai/guide/minibatch.html#chapter-6-stochastic-training-on-large-graphs
[4] https://github.com/dmlc/dgl/tree/master/examples/pytorch/GATNE-T
[5] https://networkx.github.io/documentation/latest/reference/algorithms/generated/networkx.algorithms.components.connected_components.html#networkx.algorithms.components.connected_components
[6] https://discuss.pytorch.org/t/num-workers-in-dataloader-will-increase-memory-usage/28522
[7] https://github.com/pytorch/pytorch/issues/13246
[8] https://github.com/dmlc/dgl/blob/master/examples/pytorch/pinsage/model_sparse.py
[9] https://pytorch.org/docs/stable/optim.html?highlight=sparsea#torch.optim.SparseAdam
[10] https://docs.python.org/3/library/multiprocessing.html
[11] https://github.com/dmlc/dgl/blob/master/examples/pytorch/graphsage/train_sampling_multi_gpu.py
[12] https://discuss.pytorch.org/t/should-we-split-batch-size-according-to-ngpu-per-node-when-distributeddataparallel/72769
[13] https://discuss.pytorch.org/t/distributeddataparallel-sparse-embeddings/60410/2
[14] https://github.com/pytorch/pytorch/issues/22400
[15] https://pytorch.org/docs/stable/amp.html
PaperWeekly
PaperWeekly

推荐、解读、讨论和报道人工智能前沿论文成果的学术平台。

入门图神经网络
2
相关数据
图神经网络技术

图网络即可以在社交网络或其它基于图形数据上运行的一般深度学习架构,它是一种基于图结构的广义神经网络。图网络一般是将底层图形作为计算图,并通过在整张图上传递、转换和聚合节点特征信息,从而学习神经网络基元以生成单节点嵌入向量。生成的节点嵌入向量可作为任何可微预测层的输入,并用于节点分类或预测节点之间的连接,完整的模型可以通过端到端的方式训练。

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

暂无评论
暂无评论~