Pytorch中index_select() 函數的理解

函數形式：

index_select(
    dim,
    index
)

參數：

1. dim：表示從第幾維挑選數據，類型爲int值；
2. index：表示從第一個參數維度中的哪個位置挑選數據，類型爲torch.Tensor類的實例；

功能：從張量的某個維度的指定位置選取數據。

---

代碼實例：

t = torch.arange(24).reshape(2, 3, 4) # 初始化一個tensor，從0到23，形狀爲（2,3,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,  1,  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,  5,  6,  7],
                   [ 8,  9, 10, 11]],

                  [[16, 17, 18, 19],
                   [20, 21, 22, 23]]])

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

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

---

覺得有用請留下你的贊 ^_^