GNN各阶段执行时间实验【Citeseer、Pubmed、Reddit、OGB数据集】
数据集
包括Citeseer、Pubmed、Reddit、OGB数据集的说明、数据集的划分、各数据集的介绍。
| 数据集 | 图 | 节点 | 边 | 特征(x) | 标签(y) |
|---|---|---|---|---|---|
| Citeseer | 1 | 3327 | 4732 | 3703 | 6 |
| Pubmed | 1 | 19717 | 44338 | 500 | 3 |
| 1 | 232965 | 11606919 | 602 | 41 | |
| ogbn-arxiv | 1 | 169343 | 1166243 | 128 | 40 |
| ogbn-products | 1 | 2449029 | 61859140 | 100 | 47 |
| 数据集划分（数量） | 训练集 | 验证集 | 测试集 |
|---|---|---|---|
| Citeseer | 120 | 500 | 1000 |
| Pubmed | 60 | 500 | 1000 |
| 153431 | 23831 | 55703 | |
| ogbn-arxiv | 90941 | 29799 | 48603 |
| ogbn-products | 196615 | 39323 | 2213091 |
Citeseer、Pubmed
目前有三个流行的引文网络「小数据集」：Cora、Citeseer和Pubmed。这些benchmark由Kipf和Welling在GCN模型中推出。
我的博客：节点分类任务中的引文网络benchmark。
Reddit是一个大型的在线讨论论坛，我们使用Reddit帖子构建了一个图数据集。节点代表帖子，边代表同一个用户对这两个帖子都发表了评论。这个数据集总共包含232965个节点，平均度为492。
数据集下载地址：https://data.dgl.ai/dataset/reddit.zip。
OGB数据集
因为最近主要是结合毕设，做了一些节点分类的实验，所以OGB数据集也选择是节点分类任务的ogbn-，它主要分为3类：
- 小型数据集：ogbn-arxiv
- 中型数据集：ogbn-proteins、ogbn-products、ogbn-mag
- 大型数据集：ogbn-papers100M
考虑到所做的任务和算力的限制，本次实验选择了ogbn-arxiv和ogbn-products这两个数据集。
笔记：OGB数据集。
实验结果与分析
在原先实验的基础上，增加了这几个数据集，并单独进行了实验。
代码上传到了github：https://github.com/ytchx1999/PyG-GNN-Test。
实验环境：
- GPU：Tesla T4 *1
- 显存：16G
大数据集整体上明显慢了很多。
注意，这次所测的时间计时单位仍然是ms。之所以会这样，是因为测的时间是训练阶段「每个epoch」当中各阶段的平均时间。对于小图，训练节点本身就很少，所以时间自然就短；对于大图，只能使用minibatch+sample的方法，并且每个节点只能采样少量固定的邻居，否则会导致「内存溢出」，因此时间也不会很长。
此外，PyG已经比其他的框架（DGL等）快了很多倍了，综合来看，测出来的时间并不会很长。
Citeseer
| 各阶段的执行时间/ms | Sample | Linear | Message | Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | 2.1744 | 0.0442 | 0.2120 | 0.0013 | ❎ |
| GCN（2层） | ❎ | 2.1575 | 0.0388 | 0.1399 | 0.0013 | ❎ |
| GAT（2层） | ❎ | 2.0731 | 2.7756 | 0.1316 | 0.0013 | ❎ |
| GraphSAGE（2层、minibatch） | 583.6221 | 21.9843 | 0.0188 | 11.0429 | 0.8928 | ❎ |
| JK-Nets（6层） | ❎ | 2.1501 | 0.1156 | 0.3194 | 0.0038 | 0.0864 |
Pubmed
| 各阶段的执行时间/ms | Sample | Linear | Message | Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | 2.0705 | 0.0440 | 0.2011 | 0.0013 | ❎ |
| GCN（2层） | ❎ | 2.0330 | 0.0372 | 0.1331 | 0.0012 | ❎ |
| GAT（2层） | ❎ | 2.1301 | 2.8759 | 0.1462 | 0.0012 | ❎ |
| GraphSAGE（2层、minibatch） | 586.6222 | 21.1756 | 0.0086 | 11.5174 | 0.5057 | ❎ |
| JK-Nets（6层） | ❎ | 2.1177 | 0.1143 | 0.3134 | 0.0033 | 0.0771 |
Citeseer和Pubmed的结果可以综合起来一起看，它的结果和之前进行的Cora数据集结果相似。
都是在sample和linear阶段会花费比较多的时间。因为训练数据和图本身就比较小，除去Sample，时间主要还是花在了NN的计算方面，Aggregate阶段速度比较快（图比较小导致邻居不多）。
因为使用GS方法在GPU显存不大的情况下，有些数据集会出现内存溢出RuntimeError: CUDA out of memory的情况；所以还会使用SpMM方法进行尝试（虽然也有可能会出现内存溢出的情况）
此外，在SpMM方法中，因为邻居聚合是通过稀疏矩阵乘法来实现的，所以Message+Aggregate阶段进行了合并。
最后，GS方法会涉及到采样，否则没办法运行。采样策略是，每个batch采样1024个节点，每个节点采样3-hop邻居，每hop依次为每个节点采样[15，10，5]个邻居。
使用方法：GS方法。
| 各阶段的执行时间/ms | Sample | Linear | Message | Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GCN（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GAT（2层、minibatch） | 3746.10 | 63.32 | 109.04 | 17.08 | 0.29 | ❎ |
| GraphSAGE（2层、minibatch） | 3271.21 | 47.48 | 0.35 | 48.53 | 16.84 | ❎ |
| JK-Nets（6层） | ❎ | ---- | ---- | ---- | ---- | ---- |
ogbn-arxiv
使用方法：GS方法。
| 各阶段的执行时间/ms | Sample | Linear | Message | Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GCN（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GAT（3层、minibatch） | 876.35 | 45.22 | 81.23 | 12.36 | 0.19 | ❎ |
| GraphSAGE（3层、minibatch） | 1095.46 | 36.92 | 0.26 | 37.56 | 12.87 | ❎ |
| JK-Nets（6层） | ❎ | ---- | ---- | ---- | ---- | ---- |
使用方法：SpMM方法。
因为这个数据集本身比较稠密（边比较多），所以使用稀疏矩阵方法反而在时间和空间上都完胜GS方法。因为不论是变换还是聚合都是矩阵乘法，所以Message + Aggregate阶段花费的时间要比之前长。
| 各阶段的执行时间/ms | Sample | Linear | Message + Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | ---- | ---- | ---- | ❎ |
| GCN（2层） | ❎ | 1.80 | 21.11 | 0.0018 | ❎ |
| GAT（2层） | ❎ | 1.79 | 34.84 | 0.0013 | ❎ |
| GraphSAGE（3层、minibatch） | 802.90 | 38.34 | 332.57 | 13.09 | ❎ |
| JK-Nets（6层） | ❎ | ---- | ---- | ---- | ---- |
ogbn-products
使用方法：GS方法。
| 各阶段的执行时间/ms | Sample | Linear | Message | Aggregate | Update | layer-aggregate* |
|---|---|---|---|---|---|---|
| ChebNet（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GCN（2层） | ❎ | ---- | ---- | ---- | ---- | ❎ |
| GAT（3层、minibatch） | 4016.58 | 93.32 | 196.96 | 30.33 | 0.59 | ❎ |
| GraphSAGE（3层、minibatch） | 4820.81 | 69.90 | 0.77 | 89.33 | 34.22 | ❎ |
| JK-Nets（6层） | ❎ | ---- | ---- | ---- | ---- | ---- |
这个数据集本身就比较大，使用SpMM方法+云服务器的GPU，目前还是会出现内存溢出。等换个显存大一些的机子再试一试。
从Reddit开始都可以算是大数据集了。
前面也说了，由于GPU显存本身就不大，再加上图又很大，如果不提前采样来降低计算复杂度，会导致CUDA内存溢出，根本没有办法去计算。而ChebNet和GCN模型本身没有提供采样的接口，JK-Nets内核也是GCN，所以这几个模型没有办法在这种级别的图上面运算（除非能造个几百个G的显存，或者用多GPU的分布式算法）。GAT、GraphSAGE之后的空域卷积基本上都会提供采样的接口，所以可以在OGB这样大数据集上运行。
Reddit、ogbn-arxiv、ogbn-products这三个大数据集的结果基本相似，各阶段的占比都差不多。Sample阶段大量访存，耗时最长；除此之外，GAT中因为Message阶段要计算注意力权重，所以时间也比较长；GraphSAGE中因为在大图中采样了k-hop邻居，所以Aggregate阶段在聚合的时候也需要花费比较长的时间；最后是NN计算所花费的时间。
问题与总结
在图比较大或比较稠密的时候，使用边表edge_index这种图结构数据，需要在Aggregate阶段显式的矩阵化src（x_j），这会导致很高的内存占用（甚至会导致内存溢出），反而不如使用稀疏矩阵（SparseTensor）进行存储和运算了。
初步推断，内存溢出发生在Aggregate阶段。
后面可以试一下稀疏矩阵的运算，看看能不能解决这个问题。