Tensor Dimensions

t = torch.arange(72).reshape(2,3,3,4) # 初始化一个tensor


print("t--->", t)
 
index = torch.tensor([1, 0]) # 要选取数据的位置
print("index--->", index)

data0 = t.index_select(0, index) # 第一个参数:从第0维挑选， 第二个参数:从该维中挑选的位置
print("data0--->", data1)

data1 = t.index_select(1, index)
print("data1--->", data1)
 
data2 = t.index_select(2, index) 
print("data2--->", data2)

data3 = t.index_select(3, index) 
print("data3--->", data3)

output

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

               [[24, 25, 26, 27],
                [28, 29, 30, 31],
                [32, 33, 34, 35]]],


              [[[36, 37, 38, 39],
                [40, 41, 42, 43],
                [44, 45, 46, 47]],

               [[48, 49, 50, 51],
                [52, 53, 54, 55],
                [56, 57, 58, 59]],

               [[60, 61, 62, 63],
                [64, 65, 66, 67],
                [68, 69, 70, 71]]]])
index---> tensor([1, 0])
data0---> tensor([[[[36, 37, 38, 39],
                    [40, 41, 42, 43],
                    [44, 45, 46, 47]],

                   [[48, 49, 50, 51],
                    [52, 53, 54, 55],
                    [56, 57, 58, 59]],

                   [[60, 61, 62, 63],
                    [64, 65, 66, 67],
                    [68, 69, 70, 71]]],


                  [[[ 0,  1,  2,  3],
                    [ 4,  5,  6,  7],
                    [ 8,  9, 10, 11]],

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

                   [[24, 25, 26, 27],
                    [28, 29, 30, 31],
                    [32, 33, 34, 35]]]])
data1---> tensor([[[[12, 13, 14, 15],
                    [16, 17, 18, 19],
                    [20, 21, 22, 23]],

                   [[ 0,  1,  2,  3],
                    [ 4,  5,  6,  7],
                    [ 8,  9, 10, 11]]],


                  [[[48, 49, 50, 51],
                    [52, 53, 54, 55],
                    [56, 57, 58, 59]],

                   [[36, 37, 38, 39],
                    [40, 41, 42, 43],
                    [44, 45, 46, 47]]]])
data2---> tensor([[[[ 4,  5,  6,  7],
                    [ 0,  1,  2,  3]],

                   [[16, 17, 18, 19],
                    [12, 13, 14, 15]],

                   [[28, 29, 30, 31],
                    [24, 25, 26, 27]]],


                  [[[40, 41, 42, 43],
                    [36, 37, 38, 39]],

                   [[52, 53, 54, 55],
                    [48, 49, 50, 51]],

                   [[64, 65, 66, 67],
                    [60, 61, 62, 63]]]])
data3---> tensor([[[[ 1,  0],
                    [ 5,  4],
                    [ 9,  8]],

                   [[13, 12],
                    [17, 16],
                    [21, 20]],

                   [[25, 24],
                    [29, 28],
                    [33, 32]]],


                  [[[37, 36],
                    [41, 40],
                    [45, 44]],

                   [[49, 48],
                    [53, 52],
                    [57, 56]],

                   [[61, 60],
                    [65, 64],
                    [69, 68]]]])