Преобразование модели обнаружения звука YAMNet для вывода TensorFlow Lite

Что касается машинного обучения в области акустики, большинство современных моделей основаны на распознавании шаблонов в спектрограммах с помощью сверточных нейронных сетей (CNN). В Rainforest Connection мы используем CNN для обнаружения звука бензопил, транспортных средств и других потенциально незаконных действий в тропических лесах и других охраняемых территориях. Хотя логические выводы в настоящее время выполняются на серверах в облаке, мы очень рады возможности вывода на периферии полевых датчиков для более быстрого обнаружения событий. Для этого нам нужны эффективные модели обнаружения звука, способные работать на IoT или мобильных устройствах.

YAMNet - это модель акустического обнаружения, которая классифицирует 521 различных звуков, обученных Дэном Эллисом на помеченных аудио из более чем 2 миллионов видеороликов YouTube (AudioSet). Модель YAMNet включена в исследовательский раздел репозитория моделей TensorFlow. Это потрясающий проект.

Я собираюсь пройти пошаговое руководство по преобразованию модели YAMNet в модель TensorFlow Lite, которую можно запускать на мобильных устройствах (например, развернуть на Android или iOS как Пользовательская модель Firebase ML). 🚀

1. Настройка

Загрузите Репозиторий моделей TensorFlow с GitHub и сделайте папку yamnet своим рабочим каталогом. (Я переместил его и удалил остальную часть репо, так как он не нужен для остальной части этого эксперимента.)

Получите настройку вашей любимой среды Python с помощью Tensorflow 2.2. Вы можете использовать виртуальные среды или Anaconda, если хотите, но мне проще использовать Docker:

docker run -it --rm -v ${PWD}:/app -w /app tensorflow/tensorflow:2.2.0

Далее вам необходимо скачать веса модели для YAMNet:

curl -o yamnet.h5 https://storage.googleapis.com/audioset/yamnet.h5

Затем вы можете запустить тестовый сценарий YAMNet python yamnet_test.py. Если вы видите Ran 4 tests… и OK, значит, у вас есть все, что вам нужно!

Затем давайте посмотрим, что будет на выходе при запуске модели в аудиофайле. Сценарий вывода выполнит повторную выборку файла, если он не соответствует частоте дискретизации 16 кГц модели YAMNet. Убедитесь, что у вас установлены soundfile и resampy. Если вы настраиваете свою среду с помощью Docker, просто:

apt-get install -y libsndfile1
pip install resampy soundfile

Протестируйте модель на аудиофайле:

python inference.py examples/baby_5000ms.wav

Ожидаемый результат:

examples/baby_5000ms.wav :
  Crying, sobbing: 0.555
  Baby cry, infant cry: 0.485
  Babbling    : 0.400
  Child speech, kid speaking: 0.100
  Inside, small room: 0.055

У нас есть работающая модель YAMNet на TensorFlow, оставшаяся часть этой статьи посвящена получению тех же результатов на TensorFlow Lite. 👩🏻‍💻

2. Сначала попробуйте конвертер TF Lite.

Как описано в руководстве Начало работы с TensorFlow Lite, существует несколько различных способов конвертировать модель в формат TFLite:

• Каталоги SavedModel
• Керас модели
• Конкретные функции

Глядя на yamnet.py файл, в котором определена модель, мы видим, что это модель Keras. В своей первой попытке преобразования я просто передал модель yamnet в функцию TFLiteConverter.from_keras_model следующим образом.

Когда вы запускаете python convert_fail1.py, выдает ошибки:

ValueError: None is only supported in the 1st dimension. Tensor 'input_1' has invalid shape '[1, None]'.

Урок №1: TFLite не поддерживает размеры с динамическим размером (None). Таким образом, мы не сможем передать аудиофайл произвольной длины в модель TFLite, поскольку форма входных данных должна быть определенным во время компиляции / преобразования. (Это не совсем так, поскольку TFLite добавила поддержку входов с динамическим размером в новом экспериментальном конвертере, но мне не удалось заставить его работать на мобильном устройстве). Следовательно, нам нужно выбрать длину аудиоклипа. Если посмотреть на файл params.py, частота дискретизации по умолчанию - 16,000, а длина окна - 0.96 секунд. Итак, изначально я ошибочно предположил, что форма ввода будет [1, 15360]. Чтобы объяснить, почему это неправильно, мы можем посмотреть документацию YAMNet и рассмотреть преобразование спектрограммы ...

«Спектрограмма вычисляется с использованием величин кратковременного преобразования Фурье с размером окна 25 мс, шагом окна 10 мс и периодическим окном Ханна».

Я понял это, поскольку необработанная звуковая волна имеет скользящее окно размером 0,025 с, и нам нужно достаточно окон, чтобы покрыть 0,96 с. Следовательно, длина ввода в секундах должна быть кратной 0,025 и должна быть не менее 0,96 секунды. Получается, что минимум 0,975 секунды, или 15 600 отсчетов. (Я не знаю, почему длина окна не кратна кадрам STFT или почему она не была округлена до 1 секунды - я подозреваю, что этому есть математическое объяснение!)

Чтобы указать это, мы копируем функцию yamnet_frames_model из yamnet.py и модифицируем ее так, чтобы форма ввода была зафиксирована на [1,15600].

Теперь, когда мы запускаем python convert_fail2.py, мы получаем:

error: failed while converting: ‘main’: Ops that can be supported by the flex runtime (enabled via setting the -emit-select-tf-ops flag): RFFT,ComplexAbs.

Урок №2: операторы RFFT и ComplexAbs не поддерживаются в tflite. Они фактически не используются в CNN, но являются частью генерации спектрограмм.

Затем я потратил много времени, пытаясь заставить что-то работать с SELECT_TF_OPS, добавив эту строку:

converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]

Преобразование завершилось успешно, и я думал, что мы закончили, но мне не удалось загрузить преобразованный файл TFLite на Android или iOS. (Он также не загружается в Python, потому что TensorFlow не поддерживает логический вывод с SELECT_TF_OPS, и, следовательно, локальное тестирование было невозможно - рекомендуется избегать операций выбора, если вы не знаете, что делаете.) Последовало несколько часов неудавшейся отладки!

3. Генерация спектрограмм методом TFLite.

После нескольких чашек чая я обнаружил, что решение - избегать использования неподдерживаемых операций. За это я должен поблагодарить долгий разговор по проблеме GitHub, который показал путь вперед: написание собственной TFLite-совместимой реализации tf.signal.stft. В Magenta project (Создавайте музыку и искусство с помощью машинного обучения) они сделали именно это.

В исходном коде YAMNet в строке 35 файла features.py (внутри функции waveform_to_log_mel_spectrogram) вы найдете вызов tf.signal.stft, который нам нужно заменить.

magnitude_spectrogram = tf.abs(tf.signal.stft(
        signals=waveform,
        frame_length=window_length_samples,
        frame_step=hop_length_samples,
        fft_length=fft_length))

Сейчас мы заменим его модифицированной версией функции _stft_magnitude_tflite из проекта Magenta. Скопируем только нужные нам функции в новый файл features_tflite.py. Он довольно длинный, поэтому, пожалуйста, скачайте его здесь и поместите файл в свой рабочий каталог.

Затем в нашей следующей попытке мы используем функцию waveform_log_mel_spectrogram из нашего нового файла.

python convert.py

Вуаля! Преобразование выполнено успешно, и файл yamnet.tflite сохранен. 😍🎉🕺🏻

Урок № 3: Создайте свою собственную функцию преобразования сигнала в спектрограмму на основе функции, совместимой с STFT TFLite. Спасибо, ребята, Magenta! (Забавно: после того, как я написал эту статью, я обнаружил, что функции STFT изначально были написаны самим Дэном Эллисом!)

Изучите yamnet.tflite в Netron (отличный инструмент для визуализации моделей машинного обучения) и убедитесь, что входы и выходы выглядят разумно.

• Ввод: [1,15600] (исходные аудиосэмплы продолжительностью 0,975 секунды при частоте дискретизации 16 кГц)
• Результат 0: [1,521] (прогнозы для 521 класса)
• Выход 1: [96,64] (логарифмическая спектрограмма - это может вам не понадобиться)

4. Тестирование

Вспомните сценарий inference.py, который мы использовали для тестирования модели в части 1, и давайте теперь создадим версию, которая загружает модель из сохраненного файла TFLite.

Если вы протестируете его на том же аудиофайле, который мы использовали в части 1:

python inference_tflite.py examples/baby_5000ms.wav

Тогда вы получите ошибку:

Error when checking input: expected input_1 to have shape (15600,) but got array with shape (80000,)

Помните, что ранее мы устанавливали фиксированную форму для ввода модели TFLite. Нам понадобится аудиофайл длительностью ровно 0,975 секунды. Повторите попытку с аудиофайлом правильной длины:

baby_975ms.wav :
  Crying, sobbing: 0.930
  Baby cry, infant cry: 0.862
  Whimper     : 0.157
  Babbling    : 0.112
  Inside, small room: 0.030

Запустите тот же baby_975ms.wav файл на исходной модели inference.py, и вы должны получить точно такие же результаты - мы можем быть уверены, что он работает правильно. Кроме того, у нас есть тестовый файл WAV и ожидаемые выходные данные, которые мы можем использовать для проверки его правильности работы на мобильном устройстве.

Что дальше?

Теперь у вас есть версия YAMNet для TFLite, которую вы можете добавить в приложение для Android или iOS. Я обнаружил, что самый простой способ запустить его - использовать пользовательские модели Firebase ML. Следующие руководства - хорошее место для начала.

• Использование TensorFlow Lite и ML Kit для создания пользовательских моделей машинного обучения для Android
• Пользовательские модели Firebase ML Kit для разработчиков iOS - Часть 1: Понимание TensorFlow Lite
• Пользовательские модели Firebase ML Kit для разработчиков iOS - Часть 2: Реализация Tic Tac Toe

Полный проект доступен на GitHub. Это был забавный побочный проект, позволяющий узнать больше о машинном обучении на мобильных устройствах - приветствуем ваши комментарии и предложения. 😊