How to get the trained_ddpm_checkpoint of custom data? - Githubissues

ckczzj / PDAE

Official PyTorch implementation of PDAE (NeurIPS 2022)

270 stars 20 forks source link

How to get the trained_ddpm_checkpoint of custom data? #1

Open wangxiaodong1021 opened 1 year ago

wangxiaodong1021 commented 1 year ago

How to get the trained_ddpm_checkpoint of custom data?

ckczzj commented 1 year ago

If no trained_ddpm_checkpoint exists for your data, you may need to train it by yourself using UNet. We will update the code to support regular ddpm training later.

minkyujeon commented 1 year ago

Could you share the code for regular ddpm training? I want to train this model with my custom dataset, thanks

ckczzj commented 1 year ago

We have updated the regular ddpm training code.

To train your own ddpm, you should add or modify following files:

dataset/your_dataset.py
config/your_dataset_regular.yaml
model/denoise_fn/your_dataset.py
model/denoise_fn/init.py
trainer/train_regular.py

After training, you may extract checkpoint.pt ("ema_denoise_fn") and config.yml into pre-trained-dpms folder, like other existing ones.