函数形式:

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

觉得有用请留下你的赞 ^_^

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐