根据索引沿维度选择张量切片

社会演员多 pytorch 186

原文标题Select tensor slice along a dimension based on index

我有以下形状的 PyTorch 张量:(100, 5, 100)。我需要通过从每一行中仅选择第二维中的一项来将其转换为形状为 (100, 100) 的张量,这意味着在这 5 个元素中我只需要一个,以及其对应的 100 个元素。

要执行此操作,我有第二个形状为 (100,) 的张量,其索引指定应在每行中选择这 5 个项目中的哪一个。

有没有一种简单的方法来执行此选择而不必过多地弄乱尺寸?

原文链接:https://stackoverflow.com//questions/71950956/select-tensor-slice-along-a-dimension-based-on-index

回复

我来回复
  • draw的头像
    draw 评论

    假设张量具有称为idx的标记并且具有形状(100,)。具有称为source的值的张量。然后选择:

    result = source[torch.arange(100), idx]
    
    1年前 0条评论