關于torch.scatter與torch_scatter庫的使用整理

更新時間：2023年09月11日 14:36:18 作者：回爐重造P

這篇文章主要介紹了關于torch.scatter與torch_scatter庫的使用整理,具有很好的參考價值,希望對大家有所幫助,如有錯誤或未考慮完全的地方,望不吝賜教

torch.scatter 與 tensor._scatter

Pytorch自帶的函數(shù)，用來將作為 src 的tensor根據(jù) index 的描述填充到 input 中，

形式如下：

ouput = torch.scatter(input, dim, index, src)
# 或者是
input.scatter_(dim, index, src)

兩個方法的功能是相同的，而帶下劃線的 _scatter 方法是將原tensor input 直接修改了，不帶的則會返回一個新的tensor output ， input 不變。

其中 dim 決定 index 對應值是沿著哪個維度進行修改。而 src 為數(shù)據(jù)來源，當其為tensor張量時，shape要和index相同，這樣index中每個元素都能對應 src 中對應位置的信息。

理解 scatter 方法主要是要理解 index 實現(xiàn)的 src 和 input 之間的位置對應關系，舉個例子：

dim = 0
index = torch.tensor(
	[[0, 2, 2], 
	[2, 1, 0]]
)

dim 為0時，遵循的映射原則為： input[index[i][j]][j] = src[i][j] .

也就是說，將位置 (i, j) 中 dim 對應的位置改為 index[i][j] 的值。

如位置(1,0)，index[1][0]為2，則映射后的位置為(2,0)，意味著 input 中(2,0)的位置被更改為 src 中(1,0)位置的值。

我個人形象理解是這些值會沿著dim方向滑動，上面例子中src[1][0]位置的值滑到2，成為input中的新值，這樣理解起來更形象一點。

基本理解了上面這個例子，多維情況和不同dim的情況都可以類推了。

需要注意：src和input的dtype需要相同，不然會報

Expected self.dtype to be equal to src.dtype

不一樣就先轉(zhuǎn)換再使用。

t = torch.arange(6).view(2, 3)
t = t.to(torch.float32)
print(t)
output = torch.scatter(torch.zeros((3, 3)), 0, torch.tensor([[0, 2, 2], [2, 1, 0]]), t)
print(torch.zeros((3, 3)).scatter_(0, torch.tensor([[0, 2, 2], [2, 1, 0]]), t))

輸出：

tensor([[0., 1., 2.],
[3., 4., 5.]])
tensor([[0., 0., 5.],
[0., 4., 0.],
[3., 1., 2.]])

torch_scatter庫

這個第三方庫對矩陣的分組處理這個概念做了更進一步的封裝，通過index來指定分組信息，將元素分組后進行對應處理，

最基礎的scatter方法形式如下：

torch_scatter.scatter(src, index, dim, out, dim_size, reduce)

src : 數(shù)據(jù)源
index ：分組序列
dim ：分組遵循的維度
out ：輸出的tensor，可以不指定直接讓函數(shù)輸出
dim_size ：out不指定的時候，將輸出shape變?yōu)樵撝荡笮?；dim_size也不指定，就根據(jù)計算結果來
reduce ：分組的操作，包括sum，mul，mean，min和max操作

這個方法理解關鍵在 index 的分組方法，

舉個例子：

dim = 1
index = torch.tensor([[0, 1, 1]])

torch_scatter.scatter 對 index 的順序是沒有特定規(guī)定的，相同數(shù)字對應的元素即為一組。

比如例子中，維度1上的第0個元素為一組，第1和2元素為另一組。

這樣，按照分組進行reduce定義的計算即可獲得輸出。如：

t = torch.arange(12).view(4, 3)
print(t)
t_s = torch_scatter.scatter(t, torch.tensor([[0, 1, 1]]), dim=1, reduce='sum')
print(t_s)

輸出：

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

可以看出，每行的后兩個元素求了和，與index定義相同。

要注意的是，index的 shape[0] 為1時，會自動對dim對應的維度上每一層進行相同的分組處理，如上例所示，index大小為(1, 3)，即對src的三行數(shù)據(jù)都進行了分組處理。

而另一種分組方式，如需要每行分組不同，則需要index的shape和src的shape相同，如下例：

t = torch.arange(12).view(4, 3)
print(t)
t_s = torch_scatter.scatter(t, torch.tensor([[0, 1, 1], [1, 1, 0], [0, 1, 1], [1, 1, 0]]), dim=1, reduce='sum')
print(t_s)

輸出：