Awesome
Unmasking Bias in Diffusion Model Training
<div align="center"> </div> <p align="center" style="font-size: larger;"> <a href="https://arxiv.org/abs/2310.08442">Unmasking Bias in Diffusion Model Training</a> </p> <br>What's New?
🔥 Introducing Debias: a new loss weighting strategy in noise_prediction mode that solves the bais problem in diffusion model training✨:
<p align="center"> <img src="https://github.com/yuhuUSTC/Debias/blob/main/assets/fewer_steps.png" width=85%> <p>🔥 Debias achieves both higher training efficiency and better performance with same inference steps 🚀:
<p align="center"> <img src="https://github.com/yuhuUSTC/Debias/blob/main/assets/FID_trainsample.png" width=85%> <p>Model zoo
We provide the pretrained checkpoints which can be downloaded from the following link checkpoints:
Training Scripts
You can train the base model with various targets and loss weighting:
CUDA_VISIBLE_DEVICES='2' torchrun --nproc_per_node 1 --master_port 15625 scripts/image_train.py \
--data_dir data/celeba_hq_256 \
--attention_resolutions 16 --class_cond False \
--diffusion_steps 1000 --dropout 0.0 --image_size 256 --learn_sigma True --noise_schedule linear \
--num_channels 128 --num_head_channels 64 \
--num_res_blocks 1 --resblock_updown True --use_fp16 False --use_scale_shift_norm True --lr 2e-5 --batch_size 8 \
--rescale_learned_sigmas True --p2_gamma 1 --p2_k 1 --log_dir logs/CeleAHQ/Debias \
--predict_xstart False --predict_v False \
A folder named log_dir
will be created to save the checkpoints and logs.
The default setting is noise_prediction. You can specify the training targets to x0_prediction via setting predict_xstart=True
, and v_prediction via predict_v=True
.
The different noise weighting settings in noise_prediction mode are available in L863-882 in guided_diffusion/gaussion_diffusion.py
Sampling Scripts
CUDA_VISIBLE_DEVICES='2' torchrun --nproc_per_node 1 --master_port 15625 image_sample.py \
--attention_resolutions 16 --class_cond False --diffusion_steps 1000 --dropout 0.0 --image_size 256 \
--learn_sigma True --noise_schedule linear --num_channels 128 --num_res_blocks 1 --num_head_channels 64 \
--resblock_updown True --use_fp16 False --use_scale_shift_norm True --timestep_respacing 100 \
--predict_xstart True --predict_v False \
--model_path logs/CeleAHQ/debias_ema_0.9999_500000.pt \
--sample_dir samples/CelebAHQ_100/Debias \
--sampler ddpm
You can download our pretrained checkpoint to get for model_path
. timestep_respacing
specifies the sapling steps. You can also set the sampler as DDPM, DDIM, DPM-Solver, and UniPC.
Evaluation
The evaluation follows guided_diffusion
License
This project is licensed under the MIT License - see the LICENSE file for details.
Citation
If our work assists your research, feel free to give us a star ⭐ or cite us using:
@inproceedings{Debias,
title={Unmasking Bias in Diffusion Model Training},
author={Hu Yu and Li Shen and Jie Huang and Hongsheng Li and Feng Zhao},
booktitle={The 18th European Conference on Computer Vision ECCV 2024},
year={2024},
organization={Springer}
}