Скачать BigGAN PyTorch - скачать исходный код BigGAN PyTorch

BigGAN PyTorch

Питон

1.0.0

Скачать

Biggan-Pytorch

Официально неофициальная реализация автора Биггана.

Dogball? Dogball!

Этот репо содержит код для обучения Biggans 4-8 GPU из крупномасштабной тренировки Gan для синтеза естественного изображения с высокой точностью от Эндрю Брока, Джеффа Донахью и Карен Симоньян.

Этот код от Энди Брока и Алекса Андуна.

Как использовать этот код

Вам понадобится:

Pytorch, версия 1.0.1
TQDM, Numpy, Scipy и H5Py
Набор обучения ImageNet

Во-первых, вы можете при желании подготовить предварительно обработанную версию HDF5 вашего целевого набора данных для более быстрого ввода-вывода. После этого (или нет), вам понадобятся моменты создания, необходимые для расчета FID. Они могут быть сделаны путем изменения и запуска

sh scripts/utils/prepare_data.sh

Который по умолчанию предполагает, что ваш учебный набор ImageNet загружается в data корневой папки в этом каталоге и подготовит кэшированный HDF5 с разрешением 128x128 пикселей.

В папке Scripts есть несколько сценариев Bash, которые будут обучать Biggans с разными размерами партии. Этот код предполагает, что у вас нет доступа к полному POD TPU, и, соответственно, подделывает мега-страницы, используя накопление градиента (усреднение выпускников по нескольким minibatches и выполнение только шага оптимизатора после накопления N). По умолчанию сценарий launch_BigGAN_bs256x8.sh обучает полноразмерную модель Biggan с размером партии 256 и 8 градиентных накоплений, для общего размера партии 2048 года. На 8xv100 с полной тренировкой (без тензовых ядер) этот скрипт занимает 15 дней до 150 тыс. Итераций.

Сначала вам нужно будет выяснить максимальный размер партии, который может поддержать ваша установка. Предварительно обученные модели, предоставленные здесь, прошли обучение на 8xv100 (16 ГБ VRAM каждый), что может поддерживать чуть больше, чем BS256, используемый по умолчанию. После того, как вы определите это, вы должны изменить сценарий так, чтобы размер пакета разгонялся, что количество накопления градиента равняется вашему общему размеру общего количества пакетов (по умолчанию Biggan до 2048 года).

Также обратите внимание, что в этом скрипте используется ARG --load_in_mem , который загружает весь файл (~ 64 ГБ) i128.hdf5 в ОЗУ для более быстрой загрузки данных. Если у вас недостаточно оперативной памяти, чтобы поддержать это (вероятно, 96 ГБ+), удалите этот аргумент.

Метрики и отбор проб

I believe I can fly!

Во время обучения этот скрипт будет выводить журналы с помощью обучающих метрик и метриков тестирования, сэкономит несколько копий (2 самых последних и 5 самых высоких баллов) весов/оптимизатора модели и будет производить образцы и интерполяции каждый раз, когда он спасает веса. Папка журналов содержит сценарии для обработки этих журналов и построить результаты с использованием Matlab (извините, не извините).

После обучения можно использовать sample.py для получения дополнительных выборок и интерполяций, тестировать с различными значениями усечения, размера партий, количеством скоплений постоянного статистики и т. Д. См. Сценарий sample_BigGAN_bs256x8.sh для примера.

По умолчанию все сохраняется в папках веса/образцов/журналов/данных, которые предполагаются в той же папке, что и в этом репо. Вы можете указать на все это на другую базовую папку, используя аргумент --base_root или выбрать конкретные местоположения для каждого из них с их соответствующими аргументами (например, --logs_root ).

Мы включаем сценарии, чтобы запустить Biggan-Deep, но мы не полностью обучили модель, используя их, поэтому считайте их непроверенными. Кроме того, мы включаем сценарии для запуска модели на CIFAR, а также для запуска SA-GAN (с EMA) и SN-GAN на ImageNet. Код SA-GAN предполагает, что у вас есть 4xTitanx (или эквивалентный с точки зрения оперативной памяти GPU) и будет работать с размером партии 128 и 2 градиентных накоплений.

Важное примечание по метрикам основания

В этом репо используется встроенная встроенная сеть Pytorch для расчета IS и FID. Эти оценки отличаются от результатов, которые вы получите, используя официальный код основания TF, и только для целей мониторинга! Запустите Sample.py на вашей модели, с аргументом --sample_npz , затем запустите incepation_tf13, чтобы вычислить фактический тензорфлоу. Обратите внимание, что вам потребуется установлен TensorFlow 1.3 или ранее, так как TF1.4+ Breaks Original - это код.

Предварительные модели

Pytorch Spection Score и FID Мы включаем две предварительно проведенные модельные контрольно -пропускные пункты (с G, D, копия EMA G, оптимизаторы и DICT состояния):

Основная контрольная точка для Biggan, обученного на ImageNet в 128x128, с использованием накопления градиента BS256 и 8, взятых непосредственно перед обрушением, с баллом TF 97,35 +/- 1,79: ссылка
Более ранняя контрольная точка первой модели (100 тыс.

Предварительные модели для мест-365 скоро появятся.

Этот репо также содержит сценарии для переноса оригинальных весов генератора TFHUB Biggan в Pytorch. Смотрите сценарии в папке TFHUB для получения более подробной информации.

Тонкая настройка, использование собственного набора данных или создание новых обучающих функций

That's deep, man

Если вы хотите возобновить прерванную тренировку или точно настроить предварительно обученную модель, запустите тот же сценарий запуска, но с добавленным аргументом --resume . Имена экспериментов автоматически генерируются из конфигурации, но могут быть переопределены с использованием arg --experiment_name (например, если вы хотите точно настроить модель, используя модифицированные настройки оптимизатора).

Чтобы подготовить свой собственный набор данных, вам нужно будет добавить его в наборы данных. Повторите процесс в Prepare_Data.SH (необязательно создайте предварительную копию HDF5 и рассчитайте моменты основания для FID).

По умолчанию учебный скрипт сохранит 5 лучших лучших контрольных точек, измеренных по баллу «Начало». Для наборов данных, отличных от ImageNet, оценка начала может быть очень плохим показателем качества, поэтому вы, вероятно, захотите использовать --which_best FID .

Чтобы использовать свою собственную обучающую функцию (например, обучение Bigvae): либо измените Train_fns.gan_training_function, либо добавьте новый поезд FN и добавьте его после if config['which_train_fn'] == 'GAN': line in train.py .

Аккуратный материал

Мы включаем здесь полные журналы обучения и метриков для справки. Я обнаружил, что одна из самых сложных вещей в повторном внедрении бумаги может быть проверкой, если журналы выстраиваются в очередь в начале тренировок, особенно если обучение занимает несколько недель. Надеюсь, это будет полезно для будущей работы.
Мы включаем ускоренное расчет FID-исходная версия SCIPY может потребовать более 10 минут для расчета матрицы SQRT, в этой версии используется ускоренная версия Pytorch для расчета ее за секунду.
Мы включаем ускоренную реализацию Ortho REG с низкой памяти.
По умолчанию мы вычисляем только главное значение единственного числа (спектральная норма), но этот код поддерживает вычисление большего количества SVS через аргумент --num_G_SVs .

Ключевые различия между этим кодом и оригинальным Biggan

Мы используем настройки оптимизатора от SA-GAN (G_LR = 1E-4, D_LR = 4E-4, NUM_D_STEPS = 1, в отличие от G_LR = 5E-5, D_LR = 2E-4, num_d_steps = 2). Несмотря на то, что он немного менее эффективен, это был первый угол, который мы сократили, чтобы снизить время обучения.
По умолчанию мы не используем перекрестную версию Batchnorm (он же синхронизированное пакетирование). Два варианта, которые мы попробовали (обычай, наивный, и тот, который включен в этот репо), имеют немного разные градиенты (хотя и одинаковые вперед) от встроенного паттерного штрафа, которые, по-видимому, достаточны для нанесения нанесения нанесению внимания.
Накопление градиента означает, что мы обновляем оценки SV и статистику BN в 8 раз чаще. Это означает, что статистика BN намного ближе к постоянной статистике, и что оценки единственного значения, как правило, более точны. Из -за этого мы измеряем метрики по умолчанию с G в тестовом режиме (используя пакетирование, использующую оценки статистики, вместо того, чтобы вычислять статистику постоянного уровня, как в статье). Мы все еще поддерживаем постоянную статистику (см. Сценарии Sample.sh). Это также может привести к тому, что градиенты из более ранних накоплений будут устаревшими, но на практике это не является проблемой.
Предоставленные в настоящее время предварительно проведенные модели не были обучены ортогональной регуляризации. Обучение без Ortho Reg, кажется, увеличивает вероятность того, что модели не будут поддаются усечению, но похоже, что эта конкретная модель получила выигрышный билет. Несмотря на это, мы предоставляем две высоко оптимизированные (быстрое и минимальное потребление памяти) Ortho Regal, которые непосредственно вычисляют Ortho Reg. градиенты.

Примечание о дизайне этого репо

Этот код разработан с нуля, чтобы служить расширяемой, взломанной базой для дальнейшего кода исследования. Мы задумались о том, чтобы убедиться, что абстракции являются правильной толщиной для исследований-не настолько толстые, чтобы быть непроницаемыми, но не настолько тонкими, чтобы быть бесполезными. Ключевая идея заключается в том, что если вы хотите экспериментировать с настройкой SOTA и внести некоторую модификацию (попробуйте свою собственную новую функцию потерь, архитектуру, блок самостоятельного прихода и т. Д.) Вы должны легко сделать это, просто сбросив свой код в одном или двух местах, не беспокоясь об остальной части базы кода. Такие вещи, как использование Self.which_conv и functools.partial в определении модели biggan.py, были собраны с учетом этого, как и дизайн наследства класса спектральной нормы.

С учетом сказанного, это несколько большая кодовая база для одного проекта. Хотя мы старались быть тщательными с комментариями, если есть что -то, что, по вашему мнению, может быть более ясным, лучше написанным или лучше, пожалуйста, не стесняйтесь поднять проблему или запрос на тягу.

Запросы функций

Хотите поработать или улучшить этот код? Есть несколько вещей, от которых это репо выиграет, но которые еще не работают.

Синхронизированный пакетный панель (он же перекрестный реплика). Мы пробовали два варианта этого, но по какой -то неизвестной причине он каждый раз наносит ущерб тренировке. Мы не пробовали Apex Syncbn, поскольку серверы моей школы находятся на древних драйверах Nvidia, которые не поддерживают его-Apex, вероятно, будет хорошим местом для начала.
Смешанная точная тренировка и использование тензорных ядер. Этот репо включает наивную реализацию ADAM смешанного назначения, которая работает на ранних этапах обучения, но приводит к раннему коллапсу, и ничего не делает для активации тензорных ядер (это просто снижает потребление памяти). Как указано выше, интеграция APEX в этот код и использование его методов обучения смешанного назначения, чтобы воспользоваться преимуществами тензора и снижения потребления памяти, может привести к значительному увеличению скорости.

Примечания

Смотрите этот каталог для ярлыков ImageNet.

Если вы используете этот код, пожалуйста, цитируйте

 @inproceedings{
brock2018large,
title={Large Scale {GAN} Training for High Fidelity Natural Image Synthesis},
author={Andrew Brock and Jeff Donahue and Karen Simonyan},
booktitle={International Conference on Learning Representations},
year={2019},
url={https://openreview.net/forum?id=B1xsqj09Fm},
}