549. Stable Diffusion - Image to Prompts | stable-diffusion-image-to-prompts
祝贺所有获奖者,感谢Kaggle组织了这场精彩的竞赛。同时感谢竞赛期间分享思路和代码的Kagglers们。由于我花时间进行了消融实验和延迟提交来展示各种机制的改进效果,所以略有延迟,敬请谅解。
我使用约860万条提示词生成了约1060万张图像。采用了多种方法加速图像生成过程,并使用这些图像和提示词训练了ViT-Large、ConvNeXt-XXLarge和BLIP-2模型。对这些模型进行了修改,例如移除BLIP-2的大语言模型(LLM)组件并采用LORA技术。最终提交是这三种模型嵌入向量的加权集成。
本竞赛需要多样且高质量的数据。首先使用DiffusionDB_2M数据集,优先选择与图像字幕任务相关的数据集,包括COCO Captions、VizWiz、Open Image、ADE20K、Flickr30K和TextCaps,并移除相似度>0.8的重复项。最后加入ChatGPT生成的提示词(采用更严格规则),形成约200万条高质量提示词的数据集,后文称为"PROMPTS_HQ"。每条提示词生成2张不同随机种子下的图像。
COYO-700M是包含7.47亿图像-文本对的大规模数据集。尝试从中选取部分加入PROMPT_HQ数据集但未能提升分数,因此将其用于预训练,效果显著(+0.008~0.01)。从该数据集中筛选约660万满足条件的样本(图像与文本相似度>0.3且去重),使用文本作为提示词输入Stable Diffusion生成图像,该数据集称为"PROMPTS_LQ"。
生成1060万张Stable Diffusion图像并非易事。初始脚本生成一张图约需15秒,为加速采用以下措施:
1. XFormer加速:10秒/张
2. FP16精度:5秒/张
3. 输出尺寸从768x768降至512x512,步数从50减至25:2秒/张
竞赛初期我花费数日研究稳定扩散的数学原理并设计专用模型,但未能找到高效方案,最终采用视觉模型监督训练(目标向量由all-MiniLM-L6-v2模型生成)。
最终提交包含ViT-Large、ConvNeXt-XXLarge和BLIP-2模型。虽然尝试了SWIN、DinoV2等模型,但在有限时间和资源下未显优势。
移除BLIP-2的LLM组件不仅因效率考虑,更因为本任务中提示词的语法和顺序不重要甚至有害。但若时间允许,针对提示词特性设计专用损失函数可能带来提升。
采用LORA微调模型,仅应用于底层网络并取得一定改进。
令人惊讶的是,在输出层前增加一个大型全连接层显著提升了CLIP模型性能:
ebd_dim = 1024
fc_dim = 16 * 1024
self.head = nn.Sequential(
nn.Linear(ebd_dim, fc_dim),
nn.BatchNorm1d(fc_dim),
nn.ReLU(),
nn.Linear(fc_dim, 384),
)
模型分两阶段训练:先在PROMPTS_LQ上预训练,再在PROMPTS_HQ上微调。每轮训练随机选择同一提示词的图像。尝试过不同模型训练再集成,但效果相近。
更大图像尺寸带来更高分数,但未尝试336x336以上尺寸,因计算资源应投入更重要的优化。CLIP模型训练未使用增强,BLIP-2使用默认增强。
表1展示了详细训练配置:
| 模型 | 阶段 | 轮次 | 预热 | 初始学习率 | 最终学习率 | EMA | 增强 |
|---|---|---|---|---|---|---|---|
| ViT-Large | 预训练 | 2 | 1 | 1e-4 | 8e-5 | ✓ | |
| ViT-Large | 微调 | 3 | 1 | 1e-4 | 5e-5 | ✓ | |
| ConvNeXt-XXLarge | 预训练 | 2 | 1 | 1e-4 | 8e-5 | ✓ | |
| ConvNeXt-XXLarge | 微调 | 3 | 1 | 1e-4 | 5e-5 | ✓ | |
| BLIP-2 | 预训练 | 3 | 1 | 1e-4 | 4e-5 | ✓ | ✓ |
| BLIP-2 | 微调 | 4 | 1 | 1e-4 | 2e-5 | ✓ | ✓ |
表2展示了各项策略的效果:
| 序号 | 模型 | 样本数 | 图像尺寸 | 训练参数 | LORA | 全连接层 | 私有榜 | 公开榜 |
|---|---|---|---|---|---|---|---|---|
| 1 | ViT-Large | 0.2M | 224 | 25% | 0.5579 | 0.5573 | ||
| 2 | ViT-Large | 0.2M | 336 | 25% | 0.5612 | 0.5617 | ||
| 3 | ViT-Large | 0.2M | 336 | 30% | ✓ | 0.5674 | 0.5687 | |
| 4 | ViT-Large | 0.2M | 336 | 47% | ✓ | ✓ | 0.5725 | 0.5737 |
| 5 | ViT-Large | 2M | 336 | 47% | ✓ | ✓ | 0.6368 | 0.6386 |
| 6 | ViT-Large | 6.6M+2M | 336 | 47% | ✓ | ✓ | 0.6459 | 0.6461 |
| 7 | ConvNeXt-XXLarge | 6.6M+2M | 256 | 47% | ✓ | ✓ | 0.6505 | 0.6518 |
| 8 | BLIP-2 | 6.6M+2M | 224 | 20% | ✓ | 0.6458 | 0.6489 | |
| 9 | 集成(6+7+8) | 0.6620 | 0.6634 | |||||
| 10 | 集成(11个多样预测) | 0.6632 | 0.6645 |
此笔记本包含三个模型的集成,得分0.6620/0.6634
在撰写本解决方案时,使用ChatGPT进行了校对(包括此句)。:)