pytorch中的scatter_add_函数的使用解读

pytorch scatter_add_函数的使用

关于这个函数，很少博客详细的介绍。下面就我个人理解简单介绍下。

函数：

self_tensor.scatter_add_(dim, index_tensor, other_tensor) → 输出tensor

该函数的意思是：

将other_tensor中的数据，按照index_tensor中的索引位置，添加至self_tensor矩阵中。

参数：

• dim：表示需要改变的维度，但是注意，假如dim=1，并不是说self_tensor在dim=0上的数据不会改变，这个dim只是在取矩阵数据时不固定dim=1维的索引，使用index_tensor矩阵中的索引。可能这样说还是不太理解，下面会用例子说明。其中self_tensor表示我们需要改变的tensor矩阵
• index_tensor：索引矩阵；
• other_tensor：需要添加到self_tensor中的tensor

要求：

1、self_tensor，index_tensor, other_tensor 的维度需要相同，即self.tensor.dim() = index_tensor.dim() = other_tensor.dim()；

2、假设dim=d，那么index_tensor矩阵中的所有数据必须小于d-1；

3、假设dim=d，index_tensor矩阵在d维度上的size必须小于self_tensor和other_tensor的size；即index.size(d) <= other_tensor.size(d) 且index.size(d) <= self_tensor.size(d)

三维计算公式：

self[index[i][j][k]][j][k] += other[i][j][k] # 如果 dim == 0
self[i][index[i][j][k]][k] += other[i][j][k] # 如果 dim == 1
self[i][j][index[i][j][k]] += other[i][j][k] # 如果 dim == 2

二维计算公式：

self[index[i][j]][j] += other[i][j] # 如果 dim == 0
self[i][index[i][j]] += other[i][j] # 如果 dim == 1

  index_tensor = torch.tensor([[0,1],[1,1]])
  print('index_tensor: \n', index_tensor)
  self_tensor = torch.arange(0, 4).view(2, 2)
  print('self_tensor: \n', self_tensor)
  other_tensor = torch.arange(5, 9).view(2, 2)
  print('other_tensor: \n', other_tensor)
  dim = 0
  for i in range(index_tensor.size(0)):
      for j in range(index_tensor.size(1)):
          replace_index = index_tensor[i][j]
          if dim == 0:
              # self矩阵的第0维索引
              self_tensor[replace_index][j] += other_tensor[i][j]
          elif dim == 1:
              # self矩阵的第1维索引
              self_tensor[i][replace_index] += other_tensor[i][j]       
  print(self_tensor)

结果：

    index_tensor: 
 tensor([[0, 1],
        [1, 1]])
self_tensor: 
 tensor([[0, 1],
        [2, 3]])
other_tensor: 
 tensor([[5, 6],
        [7, 8]])
tensor([[ 5,  1],
        [ 9, 17]])

使用函数计算：

index_tensor = torch.tensor([[0,1],[1,1]])
print('index_tensor: \n', index_tensor)
self_tensor = torch.arange(0, 4).view(2, 2)
print('self_tensor: \n', self_tensor)
other_tensor = torch.arange(5, 9).view(2, 2)
print('other_tensor: \n', other_tensor)
self_tensor.scatter_add_(0, index_tensor, other_tensor) 
print(self_tensor)

结果：

index_tensor: 
 tensor([[0, 1],
        [1, 1]])
self_tensor: 
 tensor([[0, 1],
        [2, 3]])
other_tensor: 
 tensor([[5, 6],
        [7, 8]])
tensor([[ 5,  1],
        [ 9, 17]])

scatter_add()函数通俗理解

self [ index[i,j] ， j ] += src [ i , j ] # if dim == 0
self [ i , index[i,j] ] += src[ i, j ] # if dim == 1

理解scatter_add()函数，看index就行了，index有多少个，self坐标就会变多少次。

self是一个二维的数组，self[第一维，第二维]，dim==0，就是将src对应坐标，对应到 index 坐标里面的值，放置到self的第一维中。

例如：

src[i,j]对应到index[i,j],假设 index[i,j] ==0,则self[第一维，第二维] 为self[0,j],只改变第一维，第二维的值和src第二维一样。

然后self[0,j]的值就会变为 self[0,j]=self[0,j]+src[i,j]

代码中 self=torch.zeros(3,5), dim=0, index=[0,1,2,0,0], src=torch.ones(2,5)

我们只看 src，当 src[0,0]=1, index[0,0]=0, self[0,0]=self[0,0]+src[0,0]=1。

当src[0,1]=1, index[0,1]=1, self[1,1]=self[1,1]+src[0,1]=1, self的第一维是index的值决定的为1，第二维是src的第二维坐标决定也为1。

当index的值没有时，就停止变换，self没有变换过的坐标值就保持不变。

总结

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。

最新评论