快捷導(dǎo)航

pytorch torch.gather函數(shù)的使用

更新時間：2024年09月09日 10:23:22 作者：qq_27390023

torch.gather 是 PyTorch 中用于在指定維度上通過索引從源張量中提取元素的函數(shù),它需要輸入張量、維度索引和索引張量,示例代碼展示了如何使用 torch.gather 從輸入張量中按索引提取元素,返回的結(jié)果張量形狀與索引張量相同

pytorch torch.gather函數(shù)

torch.gather 是 PyTorch 中的一個用于從給定維度上按索引取值的函數(shù)。

它根據(jù)一個索引張量 index，從源張量 input 中收集值，并返回一個新的張量。

torch.gather 常用于需要從張量的特定位置抽取元素的操作。

1. 函數(shù)簽名

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

input：輸入張量，表示要從中收集元素的源張量。
dim：要收集的維度索引。例如，對于一個二維張量，0 表示沿著行的維度，1 表示沿著列的維度。
index：索引張量，其形狀應(yīng)與input張量在除了dim維度之外的其他維度上保持一致。索引張量中的值表示在input張量對應(yīng)維度上要收集的元素的索引。
out（可選）：輸出張量，如果提供，結(jié)果將存儲在這個張量中。

2. 工作原理

torch.gather 在 dim 維度上，通過 index 指定的索引，從 input 中選取元素。

返回的張量的形狀與 index 的形狀相同。

3. 示例代碼

以下是一個簡單的示例代碼，演示如何使用 torch.gather 函數(shù)：

import torch

# 創(chuàng)建一個源張量
input = torch.tensor([[1, 2, 3],
                      [4, 5, 6],
                      [7, 8, 9]])

# 創(chuàng)建一個索引張量
index = torch.tensor([[0, 2, 1],
                      [2, 0, 1],
                      [1, 2, 0]])

# 在 dim=1 維度上使用 gather 函數(shù)
result = torch.gather(input, dim=1, index=index)

print("Input Tensor:")
print(input)
print("\nIndex Tensor:")
print(index)
print("\nResult Tensor:")
print(result)

4. 輸出結(jié)果

Input Tensor:
tensor([[1, 2, 3],
[4, 5, 6],
[7, 8, 9]])

Index Tensor:
tensor([[0, 2, 1],
[2, 0, 1],
[1, 2, 0]])

Result Tensor:
tensor([[1, 3, 2],
[6, 4, 5],
[8, 9, 7]])