Atlantis项目深度图到水下图像转换模型训练经验分享

Atlantis项目深度图到水下图像转换模型训练经验分享

Atlantis Atlantis: Enabling Underwater Depth Estimation with Stable Diffusion (arxiv2023) 项目地址: https://gitcode.com/gh_mirrors/atlantis7/Atlantis

项目背景

Atlantis项目是一个基于ControlNet架构的深度学习项目，专注于将普通深度图转换为水下风格的图像。该项目利用扩散模型技术，通过训练使模型能够根据输入的深度信息生成逼真的水下场景图像。

训练配置与参数设置

在实际训练过程中，硬件配置对训练效率有显著影响。根据项目经验，在NVIDIA 3090显卡上，推荐使用batch size为2或4，并设置梯度累积步数为4，这样的配置大约需要10小时完成训练。而对于显存较小的2080TI（22GB）显卡，可能需要将batch size降至1，这将显著增加训练时间至每个epoch约16小时。

关于训练步数设置，项目默认max_train_steps=30000，这实际上是保留自diffusers示例的预设值。实际训练中，通常在10000步左右就能观察到"突然收敛"(sudden convergence)现象，之后可以根据实际需求决定是否继续训练。

关键训练技巧

深度图处理注意事项

一个容易被忽视但至关重要的细节是深度图的位深设置。在使用MiDaS或Depth Anything等工具生成深度图时，必须确保将位深参数设置为1（8位），而非默认的2（16位）。这是因为Diffusers库在处理PNG图片时会调用Image.convert('RGB')，如果深度图是16位的，数值会被截断，导致训练时几乎所有深度值都变为1.0，模型将无法学习到有效的深度条件信息。

验证数据准备

验证阶段使用的数据需要特别注意：

1. 验证图像应使用与训练深度图相同格式的深度图
2. 验证提示词(prompt)可以使用训练完成后计划使用的生成提示词
3. 验证图像分辨率不宜过大，过大的验证图像可能导致显存溢出(OOM)

训练数据准备

训练数据应采用"三元组"形式，每组包含：

1. 深度图：建议使用8位PNG格式
2. 原始图像：水下场景图像
3. 元数据：包含描述性文本

常见问题与解决方案

训练不收敛问题

若遇到训练不收敛的情况，建议检查以下方面：

1. 深度图位深是否正确设置为8位
2. 验证阶段输入是否为深度图而非其他类型图像
3. batch size是否过小（可尝试增大batch size或增加梯度累积步数）
4. 训练步数是否足够（有时需要超过30000步才能收敛）

显存管理策略

针对显存不足的情况，可采取以下措施：

1. 降低batch size（最直接的方法，但会延长训练时间）
2. 调整输入图像分辨率（如从512降至448）
3. 确保验证图像分辨率适中，避免单张大图导致OOM
4. 检查训练过程中显存是否被有效释放

训练效果评估

训练过程中，loss曲线通常呈现剧烈波动状态，因此不宜仅凭loss值判断模型收敛情况。更有效的评估方法是定期（如每5000步）观察验证生成结果，检查：

1. 生成图像是否保持与输入深度图一致的布局
2. 水下效果是否逐渐明显
3. 生成质量是否随训练步数增加而提升

总结

Atlantis项目的深度图到水下图像转换模型训练需要注意多个技术细节，特别是深度图位深设置和验证数据准备。通过合理配置训练参数、仔细准备训练数据并正确解决显存问题，可以成功训练出能够根据深度信息生成逼真水下场景的ControlNet模型。训练过程中需要耐心观察验证结果，而非单纯依赖loss曲线来判断模型性能。

Atlantis Atlantis: Enabling Underwater Depth Estimation with Stable Diffusion (arxiv2023) 项目地址: https://gitcode.com/gh_mirrors/atlantis7/Atlantis