Skip to content

Latest commit

 

History

History
107 lines (75 loc) · 3.84 KB

README.md

File metadata and controls

107 lines (75 loc) · 3.84 KB
Raw

Intro

This is the training code for Diffusion-DPO. The script is adapted from the diffusers library.

Model Checkpoints

The below are initialized with StableDiffusion models and trained as described in the paper (replicable with launchers/ scripts assuming 16 GPUs, scale gradient accumulation accordingly).

StableDiffusion1.5

StableDiffusion-XL-1.0

Use this notebook to compare generations. It also has a sample of automatic quantative evaluation using PickScore.

Setup

pip install -r requirements.txt

Structure

  • launchers/ is examples of running SD1.5 or SDXL training
  • utils/ has the scoring models for evaluation or AI feedback (PickScore, HPS, Aesthetics, CLIP)
  • quick_samples.ipynb is visualizations from a pretrained model vs baseline
  • requirements.txt Basic pip requirements
  • train.py Main script, this is pretty bulky at >1000 lines, training loop starts at ~L1000 at this commit (ctrl-F "for epoch").
  • upload_model_to_hub.py Uploads a model checkpoint to HF (simple utility, current values are placeholder)

Running the training

Example SD1.5 launch

# from launchers/sd15.sh
export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATASET_NAME="yuvalkirstain/pickapic_v2"

# Effective BS will be (N_GPU * train_batch_size * gradient_accumulation_steps)
# Paper used 2048. Training takes ~24 hours / 2000 steps

accelerate launch --mixed_precision="fp16"  train.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --train_batch_size=1 \
  --dataloader_num_workers=16 \
  --gradient_accumulation_steps=1 \
  --max_train_steps=2000 \
  --lr_scheduler="constant_with_warmup" --lr_warmup_steps=500 \
  --learning_rate=1e-8 --scale_lr \
  --cache_dir="/export/share/datasets/vision_language/pick_a_pic_v2/" \
  --checkpointing_steps 500 \
  --beta_dpo 5000 \
   --output_dir="tmp-sd15"

Important Args

General

  • --pretrained_model_name_or_path what model to train/initalize from
  • --output_dir where to save/log to
  • --seed training seed (not set by default)
  • --sdxl run SDXL training
  • --sft run SFT instead of DPO

DPO

  • --beta_dpo KL-divergence parameter beta for DPO
  • --choice_model Model for AI feedback (Aesthetics, CLIP, PickScore, HPS)

Optimizers/learning rates

  • --max_train_steps How many train steps to take

  • --gradient_accumulation_steps

  • --train_batch_size see above notes in script for actual BS

  • --checkpointing_steps how often to save model

  • --gradient_checkpointing turned on automatically for SDXL

  • --learning_rate

  • --scale_lr Found this to be very helpful but isn't default in code

  • --lr_scheduler Type of LR warmup/decay. Default is linear warmup to constant

  • --lr_warmup_steps number of scheduler warmup steps

  • --use_adafactor Adafactor over Adam (lower memory, default for SDXL)

Data

  • --dataset_name if you want to switch from Pick-a-Pic
  • --cache_dir where dataset is cached locally (users will want to change this to fit their file system)
  • --resolution defaults to 512 for non-SDXL, 1024 for SDXL.
  • --random_crop and --no_hflip changes data aug
  • --dataloader_num_workers number of total dataloader workers

Citation

@misc{wallace2023diffusion,
      title={Diffusion Model Alignment Using Direct Preference Optimization}, 
      author={Bram Wallace and Meihua Dang and Rafael Rafailov and Linqi Zhou and Aaron Lou and Senthil Purushwalkam and Stefano Ermon and Caiming Xiong and Shafiq Joty and Nikhil Naik},
      year={2023},
      eprint={2311.12908},
      archivePrefix={arXiv},
      primaryClass={cs.CV}