torch_scatter.scatter()的使用方法详解

目录

  • 1. 参数
  • 2. 示例
    • 2.1 简单示例
    • 2.2 顺序问题
    • 2.3 维度问题

1. 参数

torch_scatter.scatter()的使用方法详解_第1张图片
具体来讲,scatter函数的作用就是将index中相同索引对应位置的src元素进行某种方式的操作,例如summean等,然后将这些操作结果按照索引顺序进行拼接。下面我用具体的例子来进行讲解。

2. 示例

2.1 简单示例

首先初始化src和index:

src = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # (3, 3)
index = torch.tensor([0, 0, 1], dtype=torch.int64)

接着使用scatter函数:

out = scatter(src, index, dim=0, reduce='mean')

我们观察index=[0, 0, 1],第0个位置和第1个位置都为0,第2个位置为1。也就是说,我们需要将src中第0个元素和第1个元素求平均变成一个元素,然后第2个元素求mean也就是本身为一个元素。如果index=[1, 0, 0],则意味着我们需要将src中第1个元素和第2个元素求平均变成一个元素,而第0个元素保持不变。

那么src中第几个元素到底是如何定义的呢?这就需要用到dim参数了。

dim=0意味着我们需要对src的维度0进行操作:

tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]])

即src中第0个元素为[1, 2, 3],第1个元素为[4, 5, 6],第2个元素为[7, 8, 9]

而如果dim=1,则第0个元素为[1, 4, 7],第1个元素为[2, 5, 8],第2个元素为[3, 6, 9]

因此,如果有以下代码:

src = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])  # (3, 3)
index = torch.tensor([0, 0, 1], dtype=torch.int64)
out = scatter(src, index, dim=0, reduce='mean')

那么我们就应该将src中的第0个元素为[1, 2, 3]和第1个元素为[4, 5, 6]求平均为[2.5, 3.5, 4.5],然后第2个元素[7, 8, 9]保持不变,即:

tensor([[2.5000, 3.5000, 4.5000],
        [7.0000, 8.0000, 9.0000]])

2.2 顺序问题

上面的例子中index=[0, 0, 1],最后结果是将src中第0个元素和第1个元素求平均放到了位置0,然后src中第2个元素保持不变放到了位置1。

如果index=[1, 1, 0],结果为:

tensor([[7.0000, 8.0000, 9.0000],
        [2.5000, 3.5000, 4.5000]])

可以发现,上述结果是将src中第2个元素[7, 8, 9]保持不变放到了位置0,然后将src中第0个元素[1, 2, 3]和第1个元素[4, 5, 6]求平均保持不变放到了位置1。

也就是说,无论index怎么变化,都是优先将index中0对应位置的操作结果进行放置。

2.3 维度问题

如果src的维度为(4, 3),而我们需要对dim=0操作,也就是一共有四个元素,那么index的长度应该为4,即以下操作是不合法的:

src = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])  # (4, 3)
index = torch.tensor([1, 1, 0], dtype=torch.int64)
out = scatter(src, index, dim=0, reduce='mean')
print(out)

报错为:

RuntimeError: The expanded size of the tensor (4) must match the existing size (3) at non-singleton dimension 0.  Target sizes: [4, 3].  Tensor sizes: [3, 1]

正确做法应该是:

src = torch.Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])  # (4, 3)
index = torch.tensor([1, 1, 0, 2], dtype=torch.int64)
out = scatter(src, index, dim=0, reduce='mean')
print(out)

输出为:

tensor([[ 7.0000,  8.0000,  9.0000],
        [ 2.5000,  3.5000,  4.5000],
        [10.0000, 11.0000, 12.0000]])

你可能感兴趣的:(PyTorch,torch_scatter,torch,python)