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

Posted 2023-02-16

【中文标题】如何在 .map 函数中访问张量形状？

【英文标题】：How to access Tensor shape within .map function?

【发布时间】：2020-10-04 21:11:20

【问题描述】：

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

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

它返回一个具有音频长度的张量。如果这个长度小于 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])

但我无法访问 shape 属性，因为它是动态的并且会随音频而变化。我试过使用tf.shape(audio)、audio.shape、audio.get_shape()，但我得到的形状值如None，这不允许我进行比较。

可以这样做吗？

【参考方案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,)