如何从张量中选择子张量

【问题标题】：How to select subtensor from a tensor如何从张量中选择子张量
【发布时间】：2019-11-12 15:10:45
【问题描述】：

我有一个张量A，尺寸为[N1, N2, N3/2, 2, N4, N5]。另一个张量B 是一个布尔索引张量，维度为[N1, N2, N3/2] = 1,0,0,0,1...。
现在，我想使用B 的索引获得一个维度为[N1, N2, N3/2, N4, N5] 的子张量C。
我怎么能用 pytorch 来实现呢？
请注意，我不倾向于使用 for 循环，因为它非常慢。

我在 pytorch 中查找了一些函数，但发现它们都不适合我的任务

【问题讨论】：

标签： numpy pytorch

【解决方案1】：

你说不通：当B 只选择A 的条目的部分复制到@987654327 时，C 怎么能拥有与A 相同的shape @??

假设B 有k 非零元素（即B 从N1*N2*N3/2 大小为N4*N5 在A 的矩阵中选择k 元素）然后C 可以有形状的 [k, N4, N5] 仅与 k < N1*N2*N3/2。

您可以使用nonzero() 将B 从逻辑索引转换为k 整数索引，然后使用它们来选择k 元素

C = A.view(-1, *A.shape[-2:])[B.nonzero(), ...]

【讨论】：