如何使用PyTorch的torch.gather算子进行人工智能操作示例？

本文共计498个文字，预计阅读时间需要2分钟。

目录

一、用法

二、代码示例

一、用法

torch.gather函数用于返回指定索引的Tensor元素。在PyTorch官方文档中的定义如下：

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

参数说明：- input：输入Tensor。- dim：指定从哪个维度进行索引。- index：索引Tensor，其值对应于要返回的元素的索引。- sparse_grad：如果为True，则计算梯度时使用稀疏梯度。- out：输出Tensor，如果提供，则将结果存储在此Tensor中。

二、代码示例pythonimport torch

创建一个Tensorinput_tensor=torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

创建索引Tensorindex_tensor=torch.tensor([0, 2, 1])

使用torch.gather返回指定索引的元素output_tensor=torch.gather(input_tensor, dim=0, index=index_tensor)

print(output_tensor)

目录

• 一、用法：
• 二、代码示例：

一、用法：

torch.gather 算子用于返回给定索引/下标的 Tensor 元素，在 pytorch 官网文档中的定义如下：

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

其用法等价于：

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

其中，input 是目标 Tensor ，即被搜索的 Tensor ；dim 是搜索维度（也是 Tensor ），index 是索引。

返回值类型：Tensor

二、代码示例：

概念看不懂没关系，一看代码便知用法。

a = torch.tensor([1, 5, 3, 6, 8]) b = torch.tensor([3]) # 索引为3 c = a.gather(0, b) # 输出a中第0维索引是3的元素：6 # 等价于 c=torch.gather(a,0,b) print(c) # tensor([6])

a = torch.tensor([[1.3, 2, 3, 4.5, 5], [2.0, 3, 0.3, 4.1, 2], [6, 7, 8, 9, 2], [10, 5, 0, 6, 8]]) b = torch.tensor([[1], [2], [3], [4]]) c = torch.gather(a, 1, b) # 输出a中第1维索引分别是1，2，3，4的元素：2，0.3，9，8 print(c) # tensor([[2.0000],[0.3000],[9.0000],[8.0000]])

以上就是pytorch人工智能之torch.gather算子用法示例的详细内容，更多关于pytorch算子torch.gather的资料请关注自由互联其它相关文章！