1. 程式人生 > 程式設計 >Pytorch中index_select() 函式的實現理解

Pytorch中index_select() 函式的實現理解

函式形式:

index_select(
 dim,index
)

引數:

  • dim:表示從第幾維挑選資料,型別為int值;
  • index:表示從第一個引數維度中的哪個位置挑選資料,型別為torch.Tensor類的例項;

剛開始學習pytorch,遇到了index_select(),一開始不太明白幾個引數的意思,後來查了一下資料,算是明白了一點。

a = torch.linspace(1,12,steps=12).view(3,4)
print(a)
b = torch.index_select(a,torch.tensor([0,2]))
print(b)
print(a.index_select(0,2])))
c = torch.index_select(a,1,torch.tensor([1,3]))
print(c)

先定義了一個tensor,這裡用到了linspace和view方法。

第一個引數是索引的物件,第二個引數0表示按行索引,1表示按列進行索引,第三個引數是一個tensor,就是索引的序號,比如b裡面tensor[0, 2]表示第0行和第2行,c裡面tensor[1,3]表示第1列和第3列。

輸出結果如下:

tensor([[ 1.,2.,3.,4.],
[ 5.,6.,7.,8.],
[ 9.,10.,11.,12.]])
tensor([[ 1.,12.]])
tensor([[ 2.,
[ 6.,
[10.,12.]])

功能:從張量的某個維度的指定位置選取資料。

程式碼例項:

t = torch.arange(24).reshape(2,3,4) # 初始化一個tensor,從0到23,形狀為(2,4)
print("t--->",t)
 
index = torch.tensor([1,2]) # 要選取資料的位置
print("index--->",index)
 
data1 = t.index_select(1,index) # 第一個引數:從第1維挑選, 第二個引數:從該維中挑選的位置
print("data1--->",data1)
 
data2 = t.index_select(2,index) # 第一個引數:從第2維挑選, 第二個引數:從該維中挑選的位置
print("data2--->",data2)

執行結果:

t---> tensor([[[ 0,2,3],
[ 4,5,6,7],
[ 8,9,10,11]],

[[12,13,14,15],
[16,17,18,19],
[20,21,22,23]]])

index---> tensor([1,2])

data1---> tensor([[[ 4,

[[16,23]]])

data2---> tensor([[[ 1,2],
[ 5,6],
[ 9,10]],

[[13,14],
[17,18],
[21,22]]])

以上就是本文的全部內容,希望對大家的學習有所幫助,也希望大家多多支援我們。