Этот репо содержит официальную реализацию Pytorch для вероятностных диффузионных диффузионных моделей в форме бумаги-подход к созданию негауссовых диффузионных моделей, применимых к различным неклидным коллекторам.
Андрей Охотин*, Дмитрий Мольчанов*, Владимир Аркипин, Григорий Бартош, Виктор Оханесан, Айбек Аланов, Дмитрий Ветров
Помощник: Сергей Холкин

Вероятностные модели диффузии (DDPMS) предоставляют основу для недавних прорывов при генеративном моделировании. Их марковская структура затрудняет определение DDPMS с помощью распределений, отличных от гауссовых или дискретных. В этой статье мы вводим DDPM в форме звезды (SS-DDPM). Его процесс диффузии в форме звезды позволяет нам обойти необходимость определения вероятностей перехода или вычислять постериоры. Мы устанавливаем двойственность между звездными и специфическими марковскими диффузиями для экспоненциального семейства распределений и выводим эффективные алгоритмы для обучения и отбора проб из SS-DDPMS. В случае гауссовых распределений SS-DDPM эквивалентен DDPM. Тем не менее, SS-DDPMS обеспечивает простой рецепт для проектирования диффузионных моделей с такими распределениями, как бета, фон Мизес-Fisher, Dirichlet, Wishart и другие, которые могут быть особенно полезны, когда данные находятся на ограниченном коллекторе. Мы оцениваем модель в разных настройках и находим ее конкурентоспособной даже по данным изображения, где бета-DDPM достигает результатов, сравнимых с гауссовым DDPM.
Основная логика SS-DDPM, описанная в каталоге «Lib/Diffusion». Этого может быть достаточно, если вы хотите
Также вы можете найти примеры использования SS-DDPM на геодезических и синтетических данных в каталоге «ноутбуки». Если вы хотите воспроизвести наши результаты, вы можете найти примеры выполнения команд для экспериментов на CIFAR10 и Text8.
Структура репо:
Этот репо протестирован с факелом == 1.12.0+cu113 Torchvision == 0,13.0+Cu113
git clone https://github.com/andrey-okhotin/star-shaped
cd star-shaped
pip install -r requirements.txt
# only if you don't have pytorch or your pytorch version < 1.11
pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
# only for experiments with synthetic data, otherwise you can just comment all 'import npeet'
git clone https://github.com/gregversteeg/NPEET.git && cd NPEET && pip install . && cd ../ && rm -rf NPEETЗагрузка содержания папки наборов данных - необходимо для всех трубопроводов. Эта команда может занять около 5 минут.
pip install py7zr gdown
rm -rf star-shaped/datasets
gdown --fuzzy https://drive.google.com/file/d/1ndXOmbNXR6pwoJ5qs1gVP0eAKU_RAl6E/view ? usp=sharing
py7zr x datasets.7z && rm datasets.7z && mv datasets star-shaped/datasetsЗагрузка контента в папке предварительного_моделей - не обязательно для тренировочных трубопроводов. Эта команда может занять около 3 минут.
pip install py7zr gdown
rm -rf star-shaped/pretrained_models
gdown --fuzzy https://drive.google.com/file/d/1Lebmsti31CwOFg4LYJYlWmlS7rGYQfVi/view ? usp=sharing
py7zr x pretrained_models.7z && rm pretrained_models.7z && mv pretrained_models star-shaped/pretrained_modelsДоступно для работы от Jupyter-Notebook в каталоге SS_DDPM/Notebooks. Там вы можете найти примеры обучения и отбора проб для
Доступно для работы с Bash в каталоге SS_DDPM
Запуск команды:
python lib/run_pipeline -gpu < gpu0_idx > _ < gpu1_idx > _ < gpu2_idx > -pipeline < pipeline_name > -logs_file < name_of_txt_file_to_write_execution_info > -port < available_port_for_processes_sync > . . . " other_pipeline_arguments "Пример короткого использования для работы на 3 графических процессорах на одном узле:
python lib/run_pipeline -gpu 0_1_2 -pipeline train_cifar10 -logs_file logs_train_cifar10.txt -port 8890 . . . " other_pipeline_arguments " Обучение бета-DDPM на 4 NVIDIA V100 (необходимость ~ 32 ГБ памяти графического процессора). Контрольные точки будут сохранены в каталоге "Checkpoints/train_beta_ss_cifar10". Графика потеря будет сохранена в каталоге "Результаты/Train_beta_ss_cifar10".
python lib/run_pipeline.py -gpu 0_1_2_3 -port 8900 -pipeline training_cifar10 -diffusion beta_ss -loss KL_rescaled -save_folder train_beta_ss_cifar10 -logs_file logs_training_beta_ss_cifar10.txt
cp checkpoints/training_beta_ss_cifar10/NCSNpp_episode0_epoch1050_model.pt pretrained_models/ncsnpp-cifar10_beta-ss.ptВыборка бета SS-DDPM на 2 NVIDIA V100. Результаты будут сохранены в каталоге «Результаты/SHAMPLING_BETA_SS_CIFAR10/GEDEST_SAMPLE».
python lib/run_pipeline.py -gpu 0_1 -port 8900 -pipeline sampling_cifar10 -diffusion beta_ss -num_sampling_steps 1000 -pretrained_model ncsnpp-cifar10_beta-ss.pt -num_samples 50000 -save_folder sampling_beta_ss_cifar10 -logs_file logs_sampling_beta_ss.txt
python -m pytorch_fid datasets/FID_cifar10_pack50000 results/sampling_beta_ss_cifar10/generated_samplesЕсли вы запустите те же команды, вы получите FID ~ 3.24.
Обучение категориальной SS-DDPM на 4 NVIDIA A100 (необходимо ~ 150 ГБ памяти графического процессора). Контрольные точки будут сохранены в каталоге "Checkpoints/Training_categorical_ss_text8". Графика потеря будет сохранена в каталоге "Результаты/Training_categorical_ss_text8".
python lib/run_pipeline.py -gpu 0_1_2_3 -port 8900 -pipeline training_text8 -diffusion categorical_ss -loss KL -save_folder training_categorical_ss_text8 -logs_file logs_training_categorical_ss.txt
cp checkpoints/training_categorical_ss_text8/T5Encoder_episode0_epoch2016_model.pt pretrained_models/t5base-text8_categorical-ss_fully-trained.ptОценка NLL в категориальном SS-DDPM на 3 NVIDIA A100. Результаты будут сохранены в каталоге «Результаты/nll_estimations».
python lib/run_pipeline.py -gpu 0_1_2 -port 8900 -pipeline estimating_nll_text8 -diffusion categorical_ss -pretrained_model t5base-text8_categorical-ss_fully-trained.pt -num_samples -1 -batch_size 1536 -dataset_part test -num_iwae_trajectories 1 -save_folder nll_text8_categorical-ss -logs_file logs_nll_text8_categorical_ss.txtЕсли вы запустите те же команды, вы получите nll ~ 1.61.
@ inproceedings { okhotin2023star ,
author = { Andrey Okhotin , Dmitry Molchanov , Vladimir Arkhipkin , Grigory Bartosh , Viktor Ohanesian , Aibek Alanov and Dmitry Vetrov },
title = { Star - Shaped Denoising Diffusion Probabilistic Models },
booktitle = { Advances in Neural Information Processing Systems },
volume = { 36 },
year = { 2023 }
}