从numpy ndarray中有效地采样矢量

Question

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

(B, 1, H, W)

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

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

但是如何根据＆＃34; flat＆＃34; 高效切片sIdx = np.random.choice((msk.flat>=0).nonzero()[0], replace=False, size=(k,)) ？抽样指数X？
也就是说，是否有一种有效的方法将sIdx的随机抽样与msk的切片相结合？

Answer 1

从展平的索引中获取这三个轴的其余部分的np.unravel_index，并简单地沿着这些轴索引X作为最终输出，如下所示 -

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