該倉庫包含用於紙星形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 }
}