Tensorflow 2.0：相当于numpy。 take_along_axis

Question

这是我的问题：我实现了一个简单的函数，return 将信号的峰值组织为矩阵。

@tf.function
def get_peaks(X, X_err):
    prominence = 0.9
    # X shape (B, N, 1)
    max_pooled = tf.nn.pool(X, window_shape=(20, ), pooling_type='MAX', padding='SAME') 
    maxima = tf.equal(X, max_pooled) #shape (1, N, 1)
    maxima = tf.cast(maxima, tf.float32)
    peaks = tf.squeeze(X * maxima) #shape (1, N, 1) ==> shape (N,)
    peaks_err = X_err * tf.squeeze(maxima)
    peaks_idxs, idxs = tf.math.top_k(peaks, k=2)
    return peaks_idxs, idxs

如你所见，输入的形状为(B, N, 1)，即批量样本，每个样本都是N个元素的一维向量。 returned idxs 和 peaks_idxs 一样正确，它们具有形状 (B, 2)，即批次中每个样本的两个最大值的位置（和峰值）。

问题是我还想取 idxs 对应的 peak_err。对于 numpy 我将使用：

np.take_along_axis(peaks_err, idxs, axis=1)

实际上 return 一个形状为 (B, 2) 的正确矩阵。我怎样才能对 tf 做同样的事情？我实际上尝试过使用 tf.gather:

tf.gather(peaks_err, idxs, axis=1)

但它不起作用，结果不正确，形状为 (B, B, 2)，并且有很多零。你知道我该如何解决吗？谢谢！

Answer 1

我已经解决了添加三行：

@tf.function
def get_local_maxima3(XC, SXC):
    prominence = 0.9
    # x shape (1, N, 1)
    max_pooled = tf.nn.pool(XC, window_shape=(20, ), pooling_type='MAX', padding='SAME') 
    maxima = tf.equal(XC, max_pooled) #shape (1, N, 1)
    maxima = tf.cast(maxima, tf.float32)
    peaks = tf.squeeze(XC * maxima) #shape (1, N, 1) ==> shape (N,)
    peaks_err = SXC * tf.squeeze(maxima)
    #maxima = tf.where(tf.greater(peaks, prominence)) # shape (N,)
    peaks, idxs = tf.math.top_k(peaks, k=2)

    idxs_shape = tf.shape(idxs)
    grid = tf.meshgrid(*(tf.range(idxs_shape[i]) for i in range(idxs.shape.ndims)), indexing='ij')
    index_full = tf.stack(grid[:-1] + [idxs], axis=-1)
    peaks_err = tf.gather_nd(peaks_err, index_full)
    return peaks, peaks_err

有效！如果您找到/有 smarter/faster 解决方案，我将不胜感激。

Tensorflow 2.0：相当于numpy。 take_along_axis

Tensorflow 2.0: which the equivalent of numpy. take_along_axis

python

slice

python-3.x

tensorflow2.0