在PyTorch中，gather函數(shù)是一個用于從張量（tensor）中收集特定索引位置上的元素的函數(shù)。它主要用于高級索引和從張量中提取特定信息。

定義（python）

gather函數(shù)的基本定義如下：

torch.gather(input, dim, index, out=None)

• input?(Tensor): 輸入張量。
• dim?(int): 沿其收集元素的維度。
• index?(LongTensor): 索引張量，其形狀與input在除了dim維度外的所有維度上都相同。
• out?(Tensor, optional): 輸出張量。

作用

gather函數(shù)的作用是根據(jù)index張量中的索引值，從input張量中沿著指定的dim維度收集元素。這可以用于提取張量中特定位置的值。

舉例講解

假設(shè)我們有一個形狀為(3, 3)的二維張量input，我們想要沿著第0個維度（即行的維度）收集元素。我們還需要一個索引張量index，它告訴我們從每一行中收集哪個元素。

	import torch
	# 創(chuàng)建一個形狀為 (3, 3) 的輸入張量
	input = torch.tensor([[1, 2, 3],
	[4, 5, 6],
	[7, 8, 9]])
	# 創(chuàng)建一個索引張量，它告訴我們在每一行中收集哪個元素
	# 例如，第0行收集第2個元素（值為3），第1行收集第0個元素（值為4），第2行收集第1個元素（值為8）
	index = torch.tensor([[2],
	[0],
	[1]])
	# 使用 gather 函數(shù)
	output = torch.gather(input, dim=0, index=index)
	print(output)

輸出將會是：

tensor：

	[4],
	[8]])

在這個例子中，gather函數(shù)沿著第0個維度（行）收集元素。對于每一行，它都使用index張量中對應(yīng)的索引值來確定要收集哪個元素。因此，輸出張量中的每個元素都是input張量中特定行和列的元素的組合。

注意，index張量的形狀是(3, 1)，這與input張量在除了第0個維度外的所有維度上的形狀相匹配。這是因為我們沿著第0個維度收集元素，所以其他維度的大小必須相同。