在PyTorch中,gather
函数是一个用于从张量(tensor)中收集特定索引位置上的元素的函数。它主要用于高级索引和从张量中提取特定信息。
定义(python)
gather
函数的基本定义如下:
|---|-------------------------------------------------|
| |torch.gather(input, dim, index, out=None)
|
input
(Tensor): 输入张量。dim
(int): 沿其收集元素的维度。index
(LongTensor): 索引张量,其形状与input
在除了dim
维度外的所有维度上都相同。out
(Tensor, optional): 输出张量。
作用
gather
函数的作用是根据index
张量中的索引值,从input
张量中沿着指定的dim
维度收集元素。这可以用于提取张量中特定位置的值。
举例讲解
假设我们有一个形状为(3, 3)
的二维张量input
,我们想要沿着第0个维度(即行的维度)收集元素。我们还需要一个索引张量index
,它告诉我们从每一行中收集哪个元素。
|---|-------------------------------------------------------------|
| |import torch
|
| | |
| |# 创建一个形状为 (3, 3) 的输入张量
|
| |input = torch.tensor([[1, 2, 3],
|
| |[4, 5, 6],
|
| |[7, 8, 9]])
|
| | |
| |# 创建一个索引张量,它告诉我们在每一行中收集哪个元素
|
| |# 例如,第0行收集第2个元素(值为3),第1行收集第0个元素(值为4),第2行收集第1个元素(值为8)
|
| |index = torch.tensor([[2],
|
| |[0],
|
| |[1]])
|
| | |
| |# 使用 gather 函数
|
| |output = torch.gather(input, dim=0, index=index)
|
| | |
| |print(output)
|
输出将会是:
tensor:
|---|-------------|
| |[4],
|
| |[8]])
|
在这个例子中,gather
函数沿着第0个维度(行)收集元素。对于每一行,它都使用index
张量中对应的索引值来确定要收集哪个元素。因此,输出张量中的每个元素都是input
张量中特定行和列的元素的组合。
注意,index
张量的形状是(3, 1)
,这与input
张量在除了第0个维度外的所有维度上的形状相匹配。这是因为我们沿着第0个维度收集元素,所以其他维度的大小必须相同。