日本免费高清视频-国产福利视频导航-黄色在线播放国产-天天操天天操天天操天天操|www.shdianci.com

學無先后，達者為師

網站首頁 編程語言 正文

pytorch人工智能之torch.gather算子用法示例_python

作者：ctrlA_ctrlC_ctrlV ? 更新時間： 2022-11-10 編程語言

一、用法：

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]])

原文鏈接：https://juejin.cn/post/7144580061570531358

欄目分類
最近更新