torch.index_select:通过选择索引然后去得到想要的tensor,针对比较长的tensor
torch.index_select(tensor,维度,选择的index)
代码示例:

import torch
#shape为(2,2,3)
a=torch.tensor([[[1,2,3],[4,5,6]],
                [[7,8,9],[10,11,12]]])
#选择索引0和索引2的tensor                
indices=torch.tensor([0,2])
#tensor为a,维度为2,索引为0和2
b=torch.index_select(a,2,indices)
print(b)

输出;

tensor([[[ 1,  3],
         [ 4,  6]],

        [[ 7,  9],
         [10, 12]]])
Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐