掩码扩散变换器 V2

MDTv2：更快的收敛和更强的性能

MDTv2 在图像合成方面取得了卓越的性能，例如，在ImageNet数据集上达到了新的最佳FID分数1.58，并且学习速度比之前的最佳模型DiT快10倍以上。

MDTv2比原始MDT快5倍。

简介

尽管扩散概率模型（DPMs）在图像合成方面取得了成功，但我们观察到它们通常缺乏上下文推理能力，无法学习图像中物体部分之间的关系，导致学习过程缓慢。为解决这一问题，我们提出了掩码扩散变换器（MDT），引入了一种掩码潜在建模方案，明确增强DPMs在图像中学习物体语义部分之间上下文关系的能力。

在训练过程中，MDT在潜在空间中对某些标记进行掩码。然后，设计了一个不对称的扩散变换器，在保持扩散生成过程的同时，从未掩码的标记预测被掩码的标记。我们的MDT可以从不完整的上下文输入中重建图像的完整信息，从而使其能够学习图像标记之间的相关关系。我们进一步改进了MDT，采用了更高效的宏观网络结构和训练策略，命名为MDTv2。

实验结果表明，MDTv2在图像合成方面取得了卓越的性能，例如，在ImageNet数据集上达到了新的最佳FID分数1.58，并且学习速度比之前的最佳模型DiT快10倍以上。

性能

模型	数据集	分辨率	FID-50K	Inception Score
MDT-XL/2	ImageNet	256x256	1.79	283.01
MDTv2-XL/2	ImageNet	256x256	1.58	314.73

预训练模型下载

模型托管在huggingface上，你也可以用以下方式下载：

from huggingface_hub import snapshot_download
models_path = snapshot_download("shgao/MDT-XL2")
ckpt_model_path = os.path.join(models_path, "mdt_xl2_v1_ckpt.pt")

huggingface演示在这里。

FID的新最佳记录。

设置

准备Pytorch >=2.0版本。下载并安装此仓库。

git clone https://github.com/sail-sg/MDT
cd MDT
pip install -e .

安装Adan优化器，Adan是一个强大的优化器，收敛速度比AdamW更快。(论文)

python -m pip install git+https://github.com/sail-sg/Adan.git

数据

对于标准数据集如ImageNet和CIFAR，请参考'dataset'进行准备。
使用自定义数据集时，将图像文件名更改为ClassID_ImgID.jpg，因为ADM的数据加载器从文件名中获取类别ID。

训练

在单个节点上训练（`run.sh`）。

export OPENAI_LOGDIR=output_mdtv2_s2
NUM_GPUS=8

MODEL_FLAGS="--image_size 256 --mask_ratio 0.30 --decode_layer 6 --model MDTv2_S_2"
DIFFUSION_FLAGS="--diffusion_steps 1000"
TRAIN_FLAGS="--batch_size 32"
DATA_PATH=/dataset/imagenet

python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS scripts/image_train.py --data_dir $DATA_PATH $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

在多个节点上训练（`run_ddp_master.sh`和`run_ddp_worker.sh`）。

# 主节点：
export OPENAI_LOGDIR=output_mdtv2_xl2
MODEL_FLAGS="--image_size 256 --mask_ratio 0.30 --decode_layer 4 --model MDTv2_XL_2"
DIFFUSION_FLAGS="--diffusion_steps 1000"
TRAIN_FLAGS="--batch_size 4"
DATA_PATH=/dataset/imagenet
NUM_NODE=8
GPU_PRE_NODE=8

python -m torch.distributed.launch --master_addr=$(hostname) --nnodes=$NUM_NODE --node_rank=$RANK --nproc_per_node=$GPU_PRE_NODE --master_port=$MASTER_PORT scripts/image_train.py --data_dir $DATA_PATH $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

# 工作节点：
export OPENAI_LOGDIR=output_mdtv2_xl2
MODEL_FLAGS="--image_size 256 --mask_ratio 0.30 --decode_layer 4 --model MDTv2_XL_2"
DIFFUSION_FLAGS="--diffusion_steps 1000"
TRAIN_FLAGS="--batch_size 4"
DATA_PATH=/dataset/imagenet
NUM_NODE=8
GPU_PRE_NODE=8

python -m torch.distributed.launch --master_addr=$MASTER_ADDR --nnodes=$NUM_NODE --node_rank=$RANK --nproc_per_node=$GPU_PRE_NODE --master_port=$MASTER_PORT scripts/image_train.py --data_dir $DATA_PATH $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

评估

评估代码来自ADM的TensorFlow评估套件。请按照evaluations文件夹中的说明设置评估环境。

采样和评估（`run_sample.sh`）：

MODEL_PATH=output_mdtv2_xl2/mdt_xl2_v2_ckpt.pt
export OPENAI_LOGDIR=output_mdtv2_xl2_eval
NUM_GPUS=8

echo 'CFG类条件采样：'
MODEL_FLAGS="--image_size 256 --model MDTv2_XL_2 --decode_layer 4"
DIFFUSION_FLAGS="--num_sampling_steps 250 --num_samples 50000  --cfg_cond True"
echo $MODEL_FLAGS
echo $DIFFUSION_FLAGS
echo $MODEL_PATH
python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS scripts/image_sample.py --model_path $MODEL_PATH $MODEL_FLAGS $DIFFUSION_FLAGS
echo $MODEL_FLAGS
echo $DIFFUSION_FLAGS
echo $MODEL_PATH
python evaluations/evaluator.py ../dataeval/VIRTUAL_imagenet256_labeled.npz $OPENAI_LOGDIR/samples_50000x256x256x3.npz

echo '类条件采样：'
MODEL_FLAGS="--image_size 256 --model MDTv2_XL_2 --decode_layer 4"
DIFFUSION_FLAGS="--num_sampling_steps 250 --num_samples 50000"
echo $MODEL_FLAGS
echo $DIFFUSION_FLAGS
echo $MODEL_PATH
python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS scripts/image_sample.py --model_path $MODEL_PATH $MODEL_FLAGS $DIFFUSION_FLAGS
echo $MODEL_FLAGS
echo $DIFFUSION_FLAGS
echo $MODEL_PATH
python evaluations/evaluator.py ../dataeval/VIRTUAL_imagenet256_labeled.npz $OPENAI_LOGDIR/samples_50000x256x256x3.npz

可视化

运行 infer_mdt.py 生成图像。

引用

@misc{gao2023masked,
      title={Masked Diffusion Transformer is a Strong Image Synthesizer}, 
      author={Shanghua Gao and Pan Zhou and Ming-Ming Cheng and Shuicheng Yan},
      year={2023},
      eprint={2303.14389},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

致谢

此代码库基于 DiT 和 ADM 构建。在此表示感谢！