在 Pytorch 中，如何使用 BoolTensor 蒙版在多个暗淡中切片张量？答案

【问题标题】：In Pytorch how to slice tensor across multiple dims with BoolTensor masks?在 Pytorch 中，如何使用 BoolTensor 蒙版在多个暗淡中切片张量？
【发布时间】：2021-05-26 15:35:54
【问题描述】：

我想使用 BoolTensor 索引在 Pytorch 中对多维张量进行切片。我希望索引张量保留索引为真的部分，而将索引为假的部分切掉。

我的代码是这样的

import torch
a = torch.zeros((5, 50, 5, 50))

tr_indices = torch.zeros((50), dtype=torch.bool)
tr_indices[1:50:2] = 1
val_indices = ~tr_indices

print(a[:, tr_indices].shape)
print(a[:, tr_indices, :, val_indices].shape)

我希望a[:, tr_indices, :, val_indices] 的形状为[5, 25, 5, 25]，但它返回[25, 5, 5]。结果是

torch.Size([5, 25, 5, 50])
torch.Size([25, 5, 5])

我很困惑。谁能解释一下为什么？

【问题讨论】：

标签： python pytorch advanced-indexing

【解决方案1】：

PyTorch 继承了其高级索引行为from Numpy。像这样切片两次应该可以达到您想要的输出：

a[:, tr_indices][..., val_indices]

【讨论】：

哦，谢谢！看来我应该了解更多关于 numpy 的信息。