这是我的问题：我实现了一个简单的函数，该函数返回组织为矩阵的信号峰值。
@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
如您所见，输入具有shape (B, N, 1)，即批处理样本，每个样本都是N个元素的一维向量。返回的idxs也是正确的peaks_idxs，它们的形状为（B，2），即批次中每个样品的两个最大值的位置（和峰值）。
问题是我也想拿peak_err对应的idxs。随着numpy我将使用：
np.take_along_axis(peaks_err, idxs, axis=1)
实际上返回正确的shape矩阵(B, 2)。如何使用tf做同样的事情？我实际上已经尝试使用tf.gather：
tf.gather(peaks_err, idxs, axis=1)
但它不起作用，对于形状（B，B，2）和很多零，结果不正确。你知道我该怎么解决吗？谢谢！
我解决了添加三行的问题：
@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
有用！如果您找到/有一个更聪明/更快的解决方案，我将不胜感激。
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句