HiFi-GAN
Created on December 2|Last edited on December 3
Comment
Архитектура
HiFi-GAN состоит из генератора, который преобразует исходную mel-спектрограмму до формата аудио (уменьшает количество каналов до 1 и вытягивает по "временному" измерению), а также двух видов дискриминаторов, задача которых отличать настоящее аудио от сгенерированного. Первый - Multi-Period Discriminator - смесь дискриминаторов, которые работают с сигналами разных периодов; второй - Multi-Scale Discriminator также является смесью дискриминаторов, которые смотрят уже на все аудио.
- Генератор состоит из нескольких чередующихся upsampler и MRF-блоков: upsampler (обратная свертка) отвечает непосредственно за трансформацию по временному измерению, а MRF (ResBlock с разными ядрами и dilation) нужен для выделения паттернов.
- MPD имеет сверточную структуру. Для обработки сигналов периода p временная размерность T преобразуется в матрицу размера (T / p x p) и после чего пропускается через несколько 2d сверток с ядром размера (k x 1) (такое ядро нужно, чтобы не смешивать информацию из разных сигналов).
- MSD также имеет сверточную структуру, каждый из дискриминаторов в смеси работает с аудио после AvgPooling с разными ядрами.
Реализация
В статье не очень подробно описана структура каждого из блоков, приведенные иллюстрации скорее путают (по-крайней мере генератор). В связи с этим было решено обратиться к двум другим статьям - предшественнику MelGan, где более подробно описана структура генератора и дискриминатора, аналогичного MSD, а также последовавшей статье HiFi++, там хорошая иллюстрация ResBlock из MFR. MPD был реализован по схеме из аппендикса оригинальной статьи.
После ознакомления с дополнительными статьями больших проблем с реализацией модели не возникло: генератор и дискриминаторы имеют довольно простую структуру, единственное, единственное, не очень очевидно было, что паддинги нужны во всех свертках, этот момент был честно подсмотрен в официальной реализации модели.
One batch test
В этот раз я вспомнила, что в colab можно менять файлы, из-за чего отладка обучения становилась гораздо удобнее. Это привело к двум вещам:
- Впервые у меня не 100+ отладочных коммитов
- Во время переноса изменений была допущена фатальная ошибка, из-за которой впоследствии ничего не училось и которую я два дня пыталась найти (но это спойлеры)
Я не стала дожидаться полного переобучения. На тот момент мне было достаточно того, что mel-loss стабильно падает, да и на аудио какие-то намеки на слова присутствовали. И да, меня не смутил loss дискриминатора, ушедший в 0, все-таки GAN обучаю впервые, не знала, что это плохой знак.
Когда что-то пошло не так
После всех исправлений и более-менее успешного запуска one-batch, я перешла к запуску основной модели. И... что-то пошло не так. Mel-loss почти сразу вышел на плато, спектрограммы выходили одни и те же, на аудио только шум. Генератор почему-то не учился.
GAN - вещь специфическая: небольшое отступление от авторской реализации может привести к впустую потраченным часам обучения, поэтому я пошла перечитывать статью и править детали. Спойлер - они мне не починили обучение, но многие оказались полезны после фикса основной моей ошибки. Опишу кратко, что попыталась изменить:
- Первое, что заметила - изначально взяла не тот lr scheduler, заменила на ExponentialLR (также увеличила clip градиентов, хотя в целом можно было вообще его убрать)
- Потом обнаружила, что обновляю lr каждый шаг, а не каждую эпоху - исправила
- По статье не смогла больше найти ошибок, ушла смотреть на авторскую реализацию - заметила, что у них есть инициализация весов, добавила
- Изначально использовала 0 для паддинга спектрограмм (из-за чего в аудио выходил скачок в конце), исправила на значение из приложенного MelSpectrogramConfig
- Попробовала попереставля��ь optimizer.zero_grad (мало ли)
- Продолжила смотреть на авторскую реализацию: они используют weight_norm в генераторе и усредняют выход MFR, чего не было в статье - попробавла добавить, снова мимо
- В чате подсказали, что можно добавить паддинг сразу в MelSpectrogram, тогда не будет проблем с разной длиной исходного и сгенерированного аудио
Как я упомянула, ничего из этого не починило само обучение. Приложу графики и сгенерированные примеры, но там ничего интересного нет. Из полезных исправлений я бы выделила починку lr и размера спектрограммы, а также инициализацию весов - остальное, либо потеряло актуальность после других исправлений, либо в целом не особо нужно (weight_norm и усреднение MFR).
Источник всех проблем
Главный источник всех бед - моя невнимательность) Когда я переносила изменения после отладки one_batch_test, я почему-то скопировала строчку с повторным прогоном через дискриминатор (после его обновления) не из файла в colab, где все было нормально, а из изначального файла, а именно первый прогон с detach у сгенерированного аудио. И, конечно, я этот detach не убрала и потом в упор не видела долгое время. Собственно поэтому ничего и не училось. Перед финальным запуском я также убрала лишнюю размерность у сгенерированного аудио (в чате писали, что из-за этого может mel не падать). Наконец все заработало!
Ниже представлены тестовые фразы. Получилось неидеально, можно было бы еще подержать, но если честно, то я не слышу особой разницы между аудио у HiFi_2 и HiFi_3.
Add a comment