star shaped
1.0.0
该仓库包含用于纸星形denoising扩散概率模型的官方Pytorch实现 - 创建适用于各种非欧盟歧管的非高斯扩散模型的方法。
由Andrey Okhotin*，Dmitry Molchanov*，Vladimir Arkhipkin，Grigory Bartosh，Viktor Ohanesian，Aibek Alanov，Dmitry Vetrov
助手：Sergei Kholkin
降级扩散概率模型（DDPM）为生成建模的最新突破奠定了基础。他们的马尔可夫结构使用除高斯或离散以外的其他分布定义DDPM很难。在本文中，我们介绍了星形DDPM（SS-DDPM）。它的星形扩散过程使我们能够绕过定义过渡概率或计算后期的需求。我们为指数分布家族的星形和特定的马尔可夫扩散之间建立了二元性，并得出了从SS-DDPM的训练和采样的有效算法。对于高斯分布，SS-DDPM等效于DDPM。但是，SS-DDPM提供了一个简单的配方，用于设计具有Beta，Von Mises（Fisher，Dirichlet，Wishart等）等分布的扩散模型，当数据放在约束歧管上时，这可能特别有用。我们在不同的设置中评估了该模型，并在图像数据上发现它具有竞争力，在该图像数据中，Beta SS-DDPM的结果与高斯DDPM相当。
主SS-DDPM逻辑中描述了“ LIB/扩散”目录。如果您想
另外，您可以在目录“笔记本”中找到使用SS-DDPM在地球和合成数据上的示例。如果您想重现我们的结果，可以在CIFAR10和TEXT8上找到命令执行的命令示例。
回购结构：
此存储库用Torch == 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下载预处理_MODELS文件夹的内容 - 不需要培训管道。该命令可能需要大约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可从Directory SS_DDPM/Notebooks中的Jupyter Notebook运行。在那里您可以找到培训和抽样的示例
可从目录中的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个GPU上运行的简短用法示例：
python lib/run_pipeline -gpu 0_1_2 -pipeline train_cifar10 -logs_file logs_train_cifar10.txt -port 8890 . . . " other_pipeline_arguments " 在4个NVIDIA V100上训练Beta SS-DDPM（需要〜32GB GPU内存）。检查点将保存在“检查点/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在2个NVIDIA V100上对Beta SS-DDPM进行采样。结果将保存在“结果/sampling_beta_ss_cifar10/generated_samples”目录中。
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。
训练4个NVIDIA A100上的分类SS-DDPM（需要〜150GB GPU内存）。检查点将保存在“检查点/triagner_categorical_ss_text8”目录中。损失图形将保存在“结果/triending_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在3个NVIDIA A100上估算分类SS-DDPM的NLL。结果将保存在目录“结果/nll_estimation”中。
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 }
}
