Nvidia опубликовали обновленную реализацию StyleGAN2. Обновленная StyleGAN2 выдает результаты, сравнимые с оригинальной StyleGAN2. При этом модель не требует более пары тысяч изображений для обучения, что на порядок меньше требуемого размера обучающей выборки для оригинальной StyleGAN2. Такая низкоресурсная архитектура позволяет применять GAN для новых задач, в которых сбор данных для обучающей выборки особенно осложнен. На датасете CIFAR-10 обновленная модель по FID метрике обходит предыдущий рекорд: с 5.59 на 2.42. Обучение модели быстрее в ~1.6x раз, инференс быстрее в ~1.3x раз, потребление GPU памяти в ~1.5x ниже.
Механизм адаптивной аугментации дискриминатора
Обучение генеративно-состязательных сетей (GAN) с использованием небольшого количества данных приводит к переобучению дискриминатора. Разработчики предлагают механизм адаптивной аугментации дискриминатора. Механизм позволяет стабилизировать обучение для случая ограниченной обучающей выборки. Подход не требует изменений в функционале ошибки или архитектурах сетей. Предложенный механизм можно использовать как во время обучения сети с нуля, как и при дообучении модели на другом датасете.
Ключевые изменения
Обновленная версия StyleGAN2 отличается от оригинальной версии по следующим характеристикам:
- ADA: Улучшенные результаты на датасетах с менее чем 30 тысячами изображений в обучающей выборке. State-of-the-art результаты на CIFAR-10;
- Поддержка mixed-precision: обучение быстрее в ~1.6x раз, инференс быстрее в ~1.3x раз, потребление GPU памяти в ~1.5x ниже;
- Улучшенные дефолтные гиперпараметры модели: Коробочные решения для датасетов разного разрешения и разного количества доступных GPU для обучения;
- Очищенная кодобаза: Рефакторинг и упрощение оригинального кода;
- Инструменты для работы из командной строки: Возможность воспроизвести обучение, описанное в статье, и сгенерировать видео с проекцией для любого набора изображений;
- Импорт обученной сети;
- Пайплайн аугментации изображений;
- Исправление багов в оригинальной реализации
