如何在 .map 函数中访问张量形状？

Question

我有多个长度的音频数据集，我想在 5 秒内裁剪所有音频 windows（这意味着 240000 个元素和 48000 采样率）。所以，在加载 .tfrecord 之后，我正在做：

audio, sr = tf.audio.decode_wav(image_data)

其中 returns 我是一个具有音频长度的 Tensor。如果这个长度小于 240000，我想重复音频内容直到它是 240000。所以我正在对所有音频进行处理，使用 tf.data.Dataset.map() 函数：

audio = tf.tile(audio, [5])

因为这是将我最短的音频填充到所需长度所需要的。

但为了提高效率，我只想对需要它的元素进行操作：

if audio.shape[0] < 240000:
  pad_num = tf.math.ceil(240000 / audio.shape[0]) #i.e. if the audio is 120000 long, the audio will repeat 2 times
  audio = tf.tile(audio, [pad_num])

但我无法访问形状属性，因为它是动态的并且在音频中有所不同。我试过使用 tf.shape(audio)、audio.shape、audio.get_shape()，但我得到的形状值类似于 None，这不允许我进行比较。

这可以吗？

Answer 1

您可以使用这样的函数：

import tensorflow as tf

def enforce_length(audio):
    # Target shape
    AUDIO_LEN = 240_000
    # Current shape
    current_len = tf.shape(audio)[0]
    # Compute number of necessary repetitions
    num_reps = AUDIO_LEN // current_len
    num_reps += tf.dtypes.cast((AUDIO_LEN % current_len) > 0, num_reps.dtype)
    # Do repetitions
    audio_rep = tf.tile(audio, [num_reps])
    # Trim to required size
    return audio_rep[:AUDIO_LEN]

# Test
examples = tf.data.Dataset.from_generator(lambda: iter([
    tf.zeros([100_000], tf.float32),
    tf.zeros([300_000], tf.float32),
    tf.zeros([123_456], tf.float32),
]), output_types=tf.float32, output_shapes=[None])
result = examples.map(enforce_length)
for item in result:
    print(item.shape)

输出：

(240000,)
(240000,)
(240000,)

如何在 .map 函数中访问张量形状？

How to access Tensor shape within .map function?

python

tensorflow

tf.data.dataset