Intro

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

Model Checkpoints

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

StableDiffusion1.5

Setup

Require CUDA==12.1 for sdpa acceleration

conda create -n expo python=3.10 && \
conda activate expo && \
pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu124&& \
pip install -r requirements.txt

wandb login

Structure

• launchers/ is examples of running SD1.5 training. Meanwhile, dist_bench.sh and dist_pick.sh are convinient scripts for running benchmark experiments.
• utils/ has the scoring models for evaluation or AI feedback (PickScore, HPS, Aesthetics, CLIP)
• 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").
• data_gen*.py Scripts for on-policy data generation. I provide 4 versions to adapt to different algorithms.

Running the training

Example SD1.5 launch

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

bash launchers/sd15.sh

Args for laucher script.

• MODEL_NAME: Current main agent model path.
• DATASET_DIR: Current dataset path.
• OUT_DIR: Current output path.
• EXPERIMENT_NAME: Current experiment name.
• PROMPT_PATH: Current prompt path for data generation.
• SELECTOR: Current reward model name.

Important Args for train.py

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 Reduces GPU memory usage

• --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)

Data

• --cache_dir where dataset is cached locally (users will want to change this to fit their file system)
• --resolution defaults to 512 for sd1.5
• --random_crop and --no_hflip changes data aug
• --dataloader_num_workers number of total dataloader workers

Integated training pipeline

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

bash launchers/pipe.sh

Args for pipe.sh

• start_iteration: Training starts from this iteration.
• total_iterations: Total number of iterations you want to run
• skip_first_training: Whether to skip the first training(If your training is interrupted by some reason, you can resume from the last iteration)

Citation