EDM2
1.0.0
Анализ и улучшение динамики обучения диффузионных моделей
Теро Каррас, Мика Аттала, Джаакко Лехтинен, Янн Хелстен, Тимо Айла, Самули Лейн
https://arxiv.org/abs/2312.02696
Ключевые вклад:
conda env create -f environment.yml -n edmconda activate edmЧтобы воспроизвести основные результаты нашей статьи, просто запустите:
python example.pyЭто минимальный автономный скрипт, который загружает наилучшую предварительно обученную модель для каждого набора данных и генерирует случайную сетку 8x8 изображений, используя оптимальные настройки Sampler. Ожидаемые результаты:
| Набор данных | Время выполнения | Справочное изображение |
|---|---|---|
| CIFAR-10 | ~ 6 сек | cifar10-32x32.png |
| Ffhq | ~ 28 сек | ffhq-64x64.png |
| AFHQV2 | ~ 28 сек | afhqv2-64x64.png |
| ImageNet | ~ 5 мин | imagenet-64x64.png |
Самый простой способ изучить различные стратегии выборки - это изменить example.py напрямую. Вы также можете включить предварительно обученные модели и/или предлагаемый нами пробоотборник EDM в свой собственный код, просто копируйте соответствующие биты. Обратите внимание, что определения класса для предварительно обученных моделей хранятся в самих огурцах и автоматически загружаются во время раскрытия с помощью torch_utils.persistence . Чтобы использовать модели во внешних сценариях Python, просто убедитесь, что torch_utils и dnnlib доступают через PYTHONPATH .
Docker : Вы можете запустить пример сценария, используя Docker следующим образом:
# Build the edm:latest image
docker build --tag edm:latest .
# Run the generate.py script using Docker:
docker run --gpus all -it --rm --user $( id -u ) : $( id -g )
-v ` pwd ` :/scratch --workdir /scratch -e HOME=/scratch
edm:latest
python example.py Примечание: изображение Docker требует выпуска драйвера Nvidia r520 или более поздней версии.
Вызов docker run может выглядеть пугающим, поэтому давайте распаковываем его содержимое здесь:
--gpus all -it --rm --user $(id -u):$(id -g) : со всеми включенными графическими процессорами запустите интерактивный сеанс с UID/GID текущего пользователя, чтобы избежать написания файлов Docker как root.-v `pwd`:/scratch --workdir /scratch : contrage croning hir (например, вершина этого Git Repo на вашем хост -машине), чтобы /scratch в контейнере и использовать его в качестве текущего рабочего режима.-e HOME=/scratch : укажите, где кэшировать временные файлы. ПРИМЕЧАНИЕ. Если вы хотите более мелкозернистый элемент управления, вы можете вместо этого установить DNNLIB_CACHE_DIR (для предварительно обученного кеша загрузки модели). Вы хотите, чтобы эти кеш -диры проживали на постоянных объемах, чтобы их содержимое было сохранено в нескольких вызовах docker run . Мы предоставляем предварительно обученные модели для предлагаемой наша конфигурация обучения (конфигурация F), а также базовую конфигурацию (config a):
Чтобы сгенерировать партию изображений, используя данную модель и пробоотборник, запустите:
# Generate 64 images and save them as out/*.png
python generate.py --outdir=out --seeds=0-63 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl Создание большого количества изображений может быть трудоемким; Рабочая нагрузка может быть распределена по нескольким графическим процессорам, запустив вышеуказанную команду, используя torchrun :
# Generate 1024 images using 2 GPUs
torchrun --standalone --nproc_per_node=2 generate.py --outdir=out --seeds=0-999 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl Настройки пробоотборника можно управлять с помощью параметров командной строки; См. python generate.py --help для получения дополнительной информации. Для достижения наилучших результатов мы рекомендуем использовать следующие настройки для каждого набора данных:
# For CIFAR-10 at 32x32, use deterministic sampling with 18 steps (NFE = 35)
python generate.py --outdir=out --steps=18
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# For FFHQ and AFHQv2 at 64x64, use deterministic sampling with 40 steps (NFE = 79)
python generate.py --outdir=out --steps=40
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-ffhq-64x64-uncond-vp.pkl
# For ImageNet at 64x64, use stochastic sampling with 256 steps (NFE = 511)
python generate.py --outdir=out --steps=256 --S_churn=40 --S_min=0.05 --S_max=50 --S_noise=1.003
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-imagenet-64x64-cond-adm.pkl Помимо нашего предлагаемого пробоотборника EDM, generate.py также может использоваться для воспроизведения абляций пробоотборника из раздела 3 нашей статьи. Например:
# Figure 2a, "Our reimplementation"
python generate.py --outdir=out --steps=512 --solver=euler --disc=vp --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Heun & our {t_i}"
python generate.py --outdir=out --steps=128 --solver=heun --disc=edm --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Our sigma(t) & s(t)"
python generate.py --outdir=out --steps=18 --solver=heun --disc=edm --schedule=linear --scaling=none
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl Чтобы вычислить расстояние на основе основания Fréchet (FID) для данной модели и Sampler, сначала генерируйте 50 000 случайных изображений, а затем сравните их с эталонной статистикой набора данных с использованием fid.py :
# Generate 50000 images and save them as fid-tmp/*/*.png
torchrun --standalone --nproc_per_node=1 generate.py --outdir=fid-tmp --seeds=0-49999 --subdirs
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# Calculate FID
torchrun --standalone --nproc_per_node=1 fid.py calc --images=fid-tmp
--ref=https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/cifar10-32x32.npz Обе вышеупомянутые команды могут быть параллелизированы по нескольким графическим процессорам путем настройки --nproc_per_node . Вторая команда обычно занимает 1-3 минуты на практике, но первая иногда может занять несколько часов, в зависимости от конфигурации. См. python fid.py --help для полного списка параметров.
Обратите внимание, что числовое значение FID варьируется в разных случайных семенах и очень чувствительно к количеству изображений. По умолчанию fid.py всегда будет использовать 50 000 сгенерированных изображений; Предоставление меньшего количества изображений приведет к ошибке, тогда как предоставление большего количества будет использовать случайное подмножество. Чтобы уменьшить эффект случайного изменения, мы рекомендуем повторить расчет несколько раз с разными семенами, например, --seeds=0-49999 , --seeds=50000-99999 и --seeds=100000-149999 . В нашей статье мы рассчитывали каждый FID три раза и сообщали о минимуме.
Также обратите внимание, что важно сравнить сгенерированные изображения с тем же набором данных, с которым модель была первоначально обучена. Чтобы облегчить оценку, мы предоставляем точную эталонную статистику, которая соответствует нашим предварительно обученным моделям:
Для ImageNet мы предоставляем два набора эталонных статистических данных, чтобы обеспечить сравнение яблок-яблоки: imagenet-64x64.npz следует использовать при оценке модели EDM ( edm-imagenet-64x64-cond-adm.pkl ), тогда как imagenet-64x64-baseline.npz должен использоваться при оценке модели базовой линии (модели базовой линии (модели базовой линии (базовая линия (базовая линия (базовая линия (модели базовой линии (базовая линия (модели базовой линии (базовая линия (модели базовой линии (базовая линия (модели базовой линии (базовая линия (модели базовой линии (модели базовой линии (базовая линия (модели базовой линии (базовая линия (модели базовой линии (базовая линия (модели базовой линии. baseline-imagenet-64x64-cond-adm.pkl ); Последний первоначально был обучен Дхаривал и Никол, используя несколько разных учебных данных.
Вы можете вычислить справочную статистику для ваших собственных наборов данных следующим образом:
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npz Наборы данных хранятся в том же формате, что и в стиле Stylegan: несжатые zip -архивы, содержащие несжатые файлы PNG и dataset.json файлов метаданных. Пользовательские наборы данных могут быть созданы из папки, содержащей изображения; См. python dataset_tool.py --help для получения дополнительной информации.
CIFAR-10: Загрузите версию CIFAR-10 Python и конвертируйте в ZIP Archive:
python dataset_tool.py --source=downloads/cifar10/cifar-10-python.tar.gz
--dest=datasets/cifar10-32x32.zip
python fid.py ref --data=datasets/cifar10-32x32.zip --dest=fid-refs/cifar10-32x32.npzFFHQ: Загрузите набор данных Flickr-Faces-HQ как изображения 1024x1024 и преобразовать в ZIP Archive при разрешении 64x64:
python dataset_tool.py --source=downloads/ffhq/images1024x1024
--dest=datasets/ffhq-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/ffhq-64x64.zip --dest=fid-refs/ffhq-64x64.npz AFHQV2: Загрузите обновленный набор данных Animal Faces-HQ ( afhq-v2-dataset ) и конвертируйте в ZIP Archive при разрешении 64x64:
python dataset_tool.py --source=downloads/afhqv2
--dest=datasets/afhqv2-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/afhqv2-64x64.zip --dest=fid-refs/afhqv2-64x64.npzImageNet: Загрузите задачу ImageNet Localization и преобразуйте в архив ZIP при разрешении 64x64:
python dataset_tool.py --source=downloads/imagenet/ILSVRC/Data/CLS-LOC/train
--dest=datasets/imagenet-64x64.zip --resolution=64x64 --transform=center-crop
python fid.py ref --data=datasets/imagenet-64x64.zip --dest=fid-refs/imagenet-64x64.npz Вы можете обучить новые модели с помощью train.py . Например:
# Train DDPM++ model for class-conditional CIFAR-10 using 8 GPUs
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-runs
--data=datasets/cifar10-32x32.zip --cond=1 --arch=ddpmpp Приведенный выше пример использует размер партии по умолчанию 512 изображений (контролируемые по --batch ), который равномерно делится между 8 графическими процессорами (управляемым с помощью --nproc_per_node ) для получения 64 изображений на графический процессор. Обучение больших моделям может исчерпывать память графического процессора; Лучший способ избежать этого-ограничить размер пакета за GPU, например, --batch-gpu=32 . Это использует накопление градиента, чтобы получить те же результаты, что и с использованием полных партий на GPU. См. python train.py --help для полного списка параметров.
Результаты каждого обучающего прогона сохраняются в недавно созданном каталоге, например training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 . Обучающие циклы экспортируют сетевые снимки ( network-snapshot-*.pkl ) и состояния обучения ( training-state-*.pt ) через регулярные промежутки времени (контролируется --snap и --dump ). Снимки сети могут использоваться для генерации изображений с помощью generate.py , и состояния обучения могут использоваться для возобновления обучения позже ( --resume ). Другая полезная информация записывается в log.txt и stats.jsonl . Чтобы мониторировать конвергенцию обучения, мы рекомендуем взглянуть на потерю обучения ( "Loss/loss" в stats.jsonl ), а также периодически оценивать FID для network-snapshot-*.pkl с использованием generate.py и fid.py
В следующей таблице перечислены точные обучающие конфигурации, которые мы использовали для получения наших предварительно обученных моделей:
| Модель | Графические процессоры | Время | Параметры |
|---|---|---|---|
| CIFAR10–32x32 -Cond -VP | 8xv100 | ~ 2 дня | --cond=1 --arch=ddpmpp |
| CIFAR10–32x32 -Cond -ve | 8xv100 | ~ 2 дня | --cond=1 --arch=ncsnpp |
| CIFAR10–32x32 -JOND -VP | 8xv100 | ~ 2 дня | --cond=0 --arch=ddpmpp |
| CIFAR10–32x32 -JOND -ve | 8xv100 | ~ 2 дня | --cond=0 --arch=ncsnpp |
| FFHQ -64X64 -JONCOND -VP | 8xv100 | ~ 4 дня | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| FFHQ -64X64 -JONCOND -VE | 8xv100 | ~ 4 дня | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2‑64X64 -JONCOND -VP | 8xv100 | ~ 4 дня | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| AFHQV2‑64X64 -JONCOND -VE | 8xv100 | ~ 4 дня | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| ImageNet -64x64 -Cond -ADM | 32xa100 | ~ 13 дней | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Для ImageNet-64 мы провели обучение на четырех узлах NVIDIA DGX A100, каждый из которых содержал 8 графических процессоров ампер с 80 ГБ памяти. Чтобы уменьшить требования к памяти графического процессора, мы рекомендуем либо обучать модель большим количеством графических процессоров, либо ограничить размер пакета за GPU с помощью --batch-gpu . Чтобы настроить многолетнюю обучение, пожалуйста, проконсультируйтесь с документацией Toochrun.
Copyright © 2022, Nvidia Corporation & Affiliates. Все права защищены.
Все материалы, включая исходный код и предварительно обученные модели, лицензируются в рамках международной лицензии Creative Commons Attribution-Noncommercial-Sharealik 4.0.
baseline-cifar10-32x32-uncond-vp.pkl и baseline-cifar10-32x32-uncond-ve.pkl получены из предварительно обученных моделей Ян Сонг, Джаша Сол-Дикштейн, Дидерик П. Кингма, Абхишек Кумар, Стефано Эмон и Бен Пуле. Модели первоначально были разделены по лицензии Apache 2.0.
baseline-imagenet-64x64-cond-adm.pkl получен из предварительно обученной модели Прафулла Дхаривал и Алексом Никол. Модель была первоначально разделена по лицензии MIT.
imagenet-64x64-baseline.npz получен из предварительно вычисленной эталонной статистики Прафулла Дхаривал и Алекса Никола. Статистика первоначально была разделена по лицензии MIT.
@inproceedings{Karras2022edm,
author = {Tero Karras and Miika Aittala and Timo Aila and Samuli Laine},
title = {Elucidating the Design Space of Diffusion-Based Generative Models},
booktitle = {Proc. NeurIPS},
year = {2022}
}
Это справочная реализация исследования и рассматривается как одноразовое падение кода. Таким образом, мы не принимаем внешние взносы в коде в форме запросов на притяжение.
Мы благодарим Джаакко Лехтинен, Минг-Ю Лю, Туомас Кинкаанними, Аксель Сауэр, Араш Вахдат и Янн Хелстен за обсуждения и комментарии, а также Теро Куосманен, Сэмюэль Кленберг и Янн Хелстен за поддержание нашей конфунктуру.
