numpy - 给定某些维度的索引，如何在 numpy 数组中选择值？

Question

我想获取给定索引的值，例如，

matrix = np.array([[0.1,0.2,0.7],[0.3,0.3,0.4],[0.2,0.2,0.6]])
indexes = np.array([0, 2, 1])

我期望得到的是

`array([0.1,0.4,0.2])`

一个简单的解决方案是

matrix[np.arange(3), indexes]  

但是，当输入是张量时，例如，我想选择CNN特征图中的最大值，其大小为[b，f，h，w]，索引为：

h_idx, w_idx = np.unravel_index(np.argmax(matrix.reshape(b, m, -1), -1), [h, w])

我想用上述索引选择另一个相同大小的张量 [b, f, h, w] 的值。我能做些什么来解决这个问题？好像用np.take_along_axis？

标签： numpyindexing