torch_{geometric/scatter}中一些函数的用法(softmax,scatter)

文章目录

前言
torch_geometric.utils.softmax
torch_scatter.scatter

前言

其实个人不是很建议学习这些关于图深度学习的库，我之前还用过另外一个库DGL，也是五花八门的快速实现一个GNN，但是我的感觉是还不如自己写。用他们的API是非常头疼的，有的没有官方文档。本文并不是我自己决定学习这个库而写出来的，而是因为别人的代码用了，但我不知道什么意思，所以看了一下，顺便记录下来。反正个人不推荐使用。

torch_geometric.utils.softmax

from torch_geometric.utils import softmax

这个东西和通常的softmax不一样，其会多接受一个参数，然后相同的元素之间进行softmax。例如我需要softmax[1,2,1]，并提供了另外一个参数[3,4,3]，那么其就会把3的进行内部softmax，4的进行内部softmax，从而[1,2,1]变成[0.5,1,0.5]。

如下：

al=torch.tensor([[1],[2],[1]])#shape[3,1]
a=torch.tensor([3,4,3])#shape[3]
softmax(al,a)#shape[3,1]

tensor([[0.5000],
[1.0000],
[0.5000]])

解释：为什么是[3,1]而不是[3]，这是因为通常我们都会有多头的注意力，此时可能就会是[3,num_heads]，那么其会逐头进行softmax。

torch_scatter.scatter

from torch_scatter import scatter

这个东西和上面这个思想有点类似，也是需要指定主参数（矩阵），然后额外参数。其根据这个额外参数，指定我们的主参数的哪几行是属于同一个东西的。

#主参数
a=torch.randint(0,2,(2,1,3))
a

tensor([[[1, 0, 0]],
[[0, 1, 1]]])

scatter(a,torch.tensor([1,1]),dim=0,reduce="sum")

然后torch.tensor([1,1])以及dim=0表示，a的第一行和第二行都是属于同一个东西的。然后reduce="sum"表示把同一行的东西加起来。

返回的东西比较有意思，由于我们的额外参数没有0，只有1，那么0类东西就没有，返回的时候就会返回0，1类东西有两个，并且会相加。如下：

tensor([[[0, 0, 0]],
[[1, 1, 1]]])

同理，考考大家，下面这个输出什么？

scatter(a,torch.tensor([0,1]),dim=0,reduce="sum")

输出原样a对吧，其表示只有第一行属于0类，相加，自己加自己，放到输出第一行，同理。。。