使用argmax在tensorflow中切割张量

Question

我想在tensorflow中创建一个动态丢失函数.我想计算信号FFT的能量,更具体地说,只计算最主要峰值周围的3号窗口.我无法在TF中实现,因为它会抛出很多错误,比如Stride和InvalidArgumentError (see above for traceback): Expected begin, end, and strides to be 1D equal size tensors, but got shapes [1,64], [1,64], and [1] instead.

我的代码是这样的:

self.spec = tf.fft(self.signal)
self.spec_mag = tf.complex_abs(self.spec[:,1:33])
self.argm = tf.cast(tf.argmax(self.spec_mag, 1), dtype=tf.int32)
self.frac = tf.reduce_sum(self.spec_mag[self.argm-1:self.argm+2], 1)

由于我的64位批处理和数据维度也是64,因此形状self.signal为(64,64).我希望只计算FFT的AC分量.由于信号是真正有价值的,只有一半的频谱可以完成工作.因此,形状self.spec_mag是(64,32).

该fft中的最大值位于self.argm具有形状的位置(64,1).

现在我想通过以下方法计算最大峰值周围3个元素的能量:self.spec_mag[self.argm-1:self.argm+2].

但是,当我运行代码并尝试获取值时self.frac,我会遇到多个错误.

Answer 1

Tensorflow 索引使用tf.Tensor.getitem：

此操作从张量中提取指定区域。该表示法与 NumPy 类似，但目前仅支持基本索引。这意味着当前不允许使用张量作为输入

所以使用tf.sliceandtf.strided_slice也是不可能的。

而 intf.gather indices定义了 的第一个维度的切片Tensor， intf.gather_nd定义indices了 的第一个N维度的切片Tensor，其中N = indices.shape[-1]

由于您想要 周围的 3 个值max，因此我使用列表理解手动提取第一个、第二个和第三个元素，然后是tf.stack

import tensorflow as tf

signal = tf.placeholder(shape=(64, 64), dtype=tf.complex64)
spec = tf.fft(signal)
spec_mag = tf.abs(spec[:,1:33])
argm = tf.cast(tf.argmax(spec_mag, 1), dtype=tf.int32)

frac = tf.stack([tf.gather_nd(spec,tf.transpose(tf.stack(
             [tf.range(64), argm+i]))) for i in [-1, 0, 1]])

frac = tf.reduce_sum(frac, 1)

对于行中第一个或最后一个元素的极端情况，这将失败argm，但它应该很容易解决。