从numpy ndarray有效地采样向量

Question

我有一个多维numpy数组X的shape：(B, dim, H, W)我想随机抽样kdim维向量X。
我可以从形状(B, 1, H, W)的msk获得样本索引：

sIdx = random.sample((msk.flat>=0).nonzero()[0], k) 

使用numpy的等效采样代码是：

sIdx = np.random.choice((msk.flat>=0).nonzero()[0], replace=False, size=(k,))

但是，如何根据“平坦”采样索引sIdx有效地切片X？
也就是说，有没有一种有效的方法将msk的随机抽样与X的切片结合起来？你知道吗

Answer 1

从展平的索引中用^{}获得其余三个轴的相应索引，并简单地沿着这些轴索引到X以获得最终输出，如下-

I,J,K = np.unravel_index(sIdx, (B, H, W))
out = X[I,:,J,K]