pytorch scatter_add_函數(shù)的使用

關(guān)于這個(gè)函數(shù)，很少博客詳細(xì)的介紹。下面就我個(gè)人理解簡單介紹下。

函數(shù)：

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

該函數(shù)的意思是：

將other_tensor中的數(shù)據(jù)，按照index_tensor中的索引位置，添加至self_tensor矩陣中。

參數(shù)：

• dim：表示需要改變的維度，但是注意，假如dim=1，并不是說self_tensor在dim=0上的數(shù)據(jù)不會改變，這個(gè)dim只是在取矩陣數(shù)據(jù)時(shí)不固定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、假設(shè)dim=d，那么index_tensor矩陣中的所有數(shù)據(jù)必須小于d-1；

3、假設(shè)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)

三維計(jì)算公式：

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

二維計(jì)算公式：

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)

結(jié)果：

    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]])

使用函數(shù)計(jì)算：

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)

結(jié)果：

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()函數(shù)通俗理解

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

理解scatter_add()函數(shù)，看index就行了，index有多少個(gè)，self坐標(biāo)就會變多少次。

self是一個(gè)二維的數(shù)組，self[第一維，第二維]，dim==0，就是將src對應(yīng)坐標(biāo)，對應(yīng)到 index 坐標(biāo)里面的值，放置到self的第一維中。

例如：

src[i,j]對應(yīng)到index[i,j],假設(shè) index[i,j] ==0,則self[第一維，第二維] 為self[0,j],只改變第一維，第二維的值和src第二維一樣。

然后self[0,j]的值就會變?yōu)?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，當(dāng) src[0,0]=1, index[0,0]=0, self[0,0]=self[0,0]+src[0,0]=1。

當(dāng)src[0,1]=1, index[0,1]=1, self[1,1]=self[1,1]+src[0,1]=1, self的第一維是index的值決定的為1，第二維是src的第二維坐標(biāo)決定也為1。

當(dāng)index的值沒有時(shí)，就停止變換，self沒有變換過的坐標(biāo)值就保持不變。

總結(jié)

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

最新評論