Как получить доступ к тензорной форме в .map?

Question

У меня есть набор данных аудио разной длины, и я хочу обрезать их все в 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 для формы, это не позволяет мне делать сравнение.

Можно ли это сделать?

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,)