Este repo contém a implementação oficial do Pytorch para os modelos probabilísticos de difusão em forma de estrela em forma de estrela-abordagem para criar modelos de difusão não gaussianos aplicáveis a vários coletores não-euclidianos.
Por Andrey Okhotin*, Dmitry Molchanov*, Vladimir Arkhipkin, Grigory Bartosh, Viktor Ohanesian, Aibek Alanov, Dmitry Vetrov
Assistente: Sergei Kholkin

Os modelos probabilísticos de difusão de denoising (DDPMS) fornecem a base para os recentes avanços na modelagem generativa. Sua estrutura markoviana dificulta a definição de DDPMs com outras distribuições além do gaussiano ou discreto. Neste artigo, apresentamos o DDPM em forma de estrela (SS-DDPM). Seu processo de difusão em forma de estrela nos permite ignorar a necessidade de definir as probabilidades de transição ou calcular posteriors. Estabelecemos dualidade entre difusões markovianas em forma de estrela e específicas para a família exponencial de distribuições e derivam algoritmos eficientes para treinamento e amostragem de SS-DDPMs. No caso das distribuições gaussianas, o SS-DDPM é equivalente ao DDPM. No entanto, os SS-DDPMs fornecem uma receita simples para projetar modelos de difusão com distribuições como beta, von Mises-peixe, dirichlet, wishart e outros, o que pode ser especialmente útil quando os dados estão em um coletor restrito. Avaliamos o modelo em diferentes configurações e achamos competitivo mesmo nos dados da imagem, onde o BETA SS-DDPM alcança resultados comparáveis a um DDPM gaussiano.
A lógica principal do SS-DDPM descrita no diretório "lib/difusão". Isso pode ser suficiente se você quiser
Além disso, você pode encontrar exemplos de uso do SS-DDPM em dados geodésicos e sintéticos no diretório "Notebooks". Se você deseja reproduzir nossos resultados, pode encontrar exemplos de execuções de comandos para experimentos no CIFAR10 e Text8.
Estrutura de repo:
Este repo testado com tocha == 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 NPEETDownload do conteúdo da pasta de conjuntos de dados - necessário para todos os pipelines. Este comando pode levar cerca de 5 minutos.
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/datasetsDownload do conteúdo da pasta Pretreden_models - não é necessário para o treinamento de pipelines. Este comando pode levar cerca de 3 minutos.
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_modelsDisponível para correr de Jupyter-Notebook no diretório SS_DDPM/Notebooks. Lá você pode encontrar exemplos de treinamento e amostragem para
Disponível para correr de Bash no diretório SS_DDPM
Comando em execução:
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 "Exemplo de uso curto para executar em 3 GPUs em um único nó:
python lib/run_pipeline -gpu 0_1_2 -pipeline train_cifar10 -logs_file logs_train_cifar10.txt -port 8890 . . . " other_pipeline_arguments " Treinamento beta SS-DDPM em 4 NVIDIA V100 (necessidade de ~ 32 GB de memória GPU). Os pontos de verificação serão salvos no diretório "pontos de verificação/trens_beta_ss_cifar10". Os gráficos de perda serão salvos no diretório "Resultados/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.ptAmostragem beta ss-ddpm em 2 nvidia v100. Os resultados serão salvos no diretório "Resultados/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_samplesSe você executar exatamente os mesmos comandos, obterá FID ~ 3.24.
SS-DDPM categórico de treinamento em 4 NVIDIA A100 (necessidade de ~ 150 GB de memória GPU). Os pontos de verificação serão salvos no diretório "pontos de verificação/treinamento_categorical_ss_text8". Os gráficos de perda serão salvos no diretório "Resultados/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.ptEstimando a NLL em SS-DDPM categórico em 3 NVIDIA A100. Os resultados serão salvos no diretório "resultados/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.txtSe você executar exatamente os mesmos comandos, obterá 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 }
}