微软工程师用PyTorch实现图注意力网络,可视化效果惊艳

2021-02-23 12:02:58 浏览数 (1)

机器之心报道

作者:陈萍、蛋酱

近日,一个关于图注意力网络可视化的项目吸引了大批研究人员的兴趣,上线仅仅一天,收获 200 星。该项目是关于用 PyTorch 实现的图注意力网络(GAT),包括易于理解的可视化。

项目地址:https://github.com/gordicaleksa/pytorch-GAT

在正式介绍项目之前,先提一下图神经网络(GNN)。GNN 是一类基于深度学习的处理图域信息的方法。由于其较好的性能和可解释性,GNN 最近已成为一种广泛应用的图分析方法。现已广泛应用于计算生物学、计算药理学、推荐系统等。

GNN 把深度学习应用到图结构 (Graph) 中,其中的图卷积网络 GCN 可以在 Graph 上进行卷积操作,但是 GCN 存在一些缺陷。因此,Bengio 团队在三年前提出了图注意力网络(GAT,Graph Attention Network) ,旨在解决 GCN 存在的问题。

GAT 是空间(卷积)GNN 的代表。由于 CNNs 在计算机视觉领域取得了巨大的成功,研究人员决定将其推广到图形上,因此 GAT 应运而生。

现在,有人用 PyTorch 实现了 GAT 可视化。我们来看看该项目是如何实现的。

可视化

Cora 可视化

说到 GNN,就不得不介绍一下 Cora 数据集。Cora 数据集由许多机器学习论文组成,是近年来图深度学习很喜欢使用的数据集。Cora 中的节点代表研究论文,链接是这些论文之间的引用。项目作者添加了一个用于可视化 Cora 和进行基本网络分析的实用程序。Cora 如下图所示:

节点大小对应于其等级(即进出边的数量)。边的粗细大致对应于边的「popular」或「连接」程度。以下是显示 Cora 上等级(进出边的数量)分布的图:

进和出的等级图是一样的,因为处理的是无向图。在底部的图(等级分布)上,我们可以看到一个有趣的峰值发生在 [2,4] 范围内。这意味着多数节点有少量的边,但是有 1 个节点有 169 条边(绿色大节点)。

注意力可视化

有了一个训练好的 GAT 模型以后,我们就可以将某些节点所学的注意力可视化。节点利用注意力来决定如何聚合周围的节点,如下图所示:

这是 Cora 节点中边数最多的节点之一(引用)。颜色表示同一类的节点。

熵直方图

另一种理解 GAT 没有在 Cora 上学习注意力模式 (即它在学习常量注意力) 的方法是,将节点邻域的注意力权重视为概率分布,计算熵,并在每个节点邻域积累信息。

我们希望 GAT 的注意力分布有偏差。你可以看到橙色的直方图是理想均匀分布的样子,而浅蓝色的是学习后的分布,它们是完全一样的。

分析 Cora 嵌入空间 (t-SNE)

GAT 的输出张量为 shape=(2708,7),其中 2708 是 Cora 中的节点数,7 是类数。用 t-SNE 把这些 7 维向量投影成 2D,得到:

使用方法

方法 1:Jupyter Notebook

只需从 Anaconda 控制台运行 Jupyter Notebook,它将在你的默认浏览器中打开 session。打开 The Annotated GAT.ipynb 即可开始。

注意,如果你得到了 DLL load failed while importing win32api: The specified module could not be found,只需要 pip uninstall pywin32,或者 pip install pywin32、onda install pywin32。

方法 2:使用你选择的 IDE

如果使用自己选择的 IDE,只需要将 Python 环境和设置部分连接起来。

训练 GAT

在 Cora 上训练 GAT 所需的一切都已经设置好了,运行时只需调用 python training_script.py

此外,你还可以:

  • 添加 --should_visualize - 以可视化你的图形数据
  • 在数据的测试部分添加 --should_test - 以评估 GAT
  • 添加 --enable_tensorboard - 开始保存度量标准(准确率、损失)

代码部分的注释很完善,因此你可以了解到训练本身是如何运行的。

该脚本将:

  • 将 checkpoint* .pth 模型转储到 models/checkpoints/
  • 将 final* .pth 模型转储到 models/binaries/
  • 将度量标准保存到中 runs/,只需 tensorboard --logdir=runs 在 Anaconda 中运行即可将其可视化
  • 定期将一些训练元数据写入控制台

通过 tensorboard --logdir=runs 在控制台中调用,并将 http://localhost:6006/URL 粘贴到浏览器中,可以在训练过程中将度量标准可视化:

可视化工具

如果要可视化 t-SNE 嵌入,请注意或嵌入该 visualize_gat_properties 函数的注释,并设置 visualization_type 为:

  • VisualizationType.ATTENTION - 如果希望可视化节点附近的注意力
  • VisualizationType.EMBEDDING - 如果希望可视化嵌入(通过 t-SNE)
  • VisualizationType.ENTROPY - 如果想可视化熵直方图

然后,你就得到了一张优秀的可视化效果图(VisualizationType.ATTENTION 可选):

硬件需求

GAT 不需要那种很强的硬件资源,尤其是如果你只想运行 Cora 的话,有 2GB 以上的 GPU 就可以了。

  • 在 RTX 2080 GPU 上训练 GAT 大约需要 10 秒;
  • 保留 1.5 GB 的 VRAM 内存(PyTorch 的缓存开销,为实际张量分配的内存少得多);
  • 模型本身只有 365 KB。

如果你想了解更多关于 GAT 的内容,请点击下方视频:

0 人点赞