4月26日，AI初创公司Mosaic ML说：

“ MosaiCML非常适合训练扩散模型，并且与以前的工具相比已取得了很大的改进。” —— Dream3D首席执行官Tony Francis

几个月前，我们演示了如何以低价在Mosiacml平台上从头开始训练大型分区模型。

图1：AI想象的菌丝形式。这种将单数图像引入设计过程的尝试有望突破创造力的界限。以上所有图都是由我们在MosaiCML平台上从头开始训练的内部扩散模型创建的。

使用您自己的数据来训练自己的图像生成模型，不久前的这个目标现在是可行的。通过培训我们自己的扩散模型，我们可以：

使用专有数据；调整某些艺术或摄影风格的表达；避免违反知识产权法，并确保可以将模型用于商业目的。我们已经开了开源的代码和方法来训练扩散模型，使您可以随意训练自己的模型（https://github.com/mosaicml/diffusion）。

设置

图2：使用创造力并拥抱意外发现。我们的扩散模型能够生成不同主题，艺术和摄影样式的图片。

模型：我们的扩散模型是由变量自动编码器（VAE），剪辑模型，U-NET和扩散器噪声调度程序组成的作曲家模型。所有功能组件均来自Huggingface的扩散器库。所有模型配置均基于稳定性/稳定- 扩散2键。

图3：扩散模型的简化说明。

数据：我们使用的培训数据集是LAION-5B的子集，其中包括具有纯英文标题的样品，其美学得分为4.5+。与稳定的扩散2基础类似，我们根据训练数据的图像分辨率将训练过程分为两个阶段。在第一阶段，我们使用了分辨率大于或等于256 x 256的图像，总计7.9亿张图像- 标题样本。在第二阶段，我们仅使用分辨率大于或等于512 x 512的图像，总计3亿个图标- 标题样本。

计算：两个训练阶段在128个NVIDIA A100 GPU上运行。第一个训练阶段花费了1.6天，总共进行了550,000个迭代。第二阶段花了4.9天，总共进行了850,000个迭代，总训练时间为20,051 A100小时。除了训练时间外，我们还预先计算了VAE和剪辑模型的潜在空间，希望减少数据集多个通行证的训练时间和成本。潜在空间计算需要大约3784 A100小时，因此该模型的总训练时间为23835 A100小时。假设A100为每小时2美元，总价为47,700美元。

技术堆栈：我们将作曲家用作培训框架，流媒体数据集加载100 TB的培训数据，并在培训和评估基础架构时使用Mosaicml平台来解决部署和管理挑战。

图4：训练期间的损失曲线。我们的平台发现了两个硬件故障，并在没有任何人干预的情况下自动完成了重新启动。损耗曲线是不连续的，因为第二阶段将分辨率从上一阶段的256 x 256增加到512 x 512。

挑战与解决方案

无论是扩散扩散模型还是大型语言模型，规模训练都需要一系列重大挑战。我们使用MosaiCML平台进行扩散模型训练，该培训会自动解决大多数问题，以确保我们可以专注于训练最佳模型。以下是规模训练中的三个主要挑战，以及我们的平台如何解决。

基础设施

大规模数据集的大型模型无疑需要巨大的计算能力。 MosaiCML平台可以轻松地在任何云服务提供商中精心编排数百个GPU。例如，我们的主要培训工作在128 A100 GPU的集群中运行。为了确保评估模型不会减慢训练，我们使用不同的云提供商自动在不同群集中的各种检查点上进行运行评估，并根据可用性将运行尺寸缩小到64甚至至少8 GPU。

即使开始培训后，软件或硬件故障也可能导致训练中断，这需要监视24/7。幸运的是，MosaiCML平台的节点医生和看门狗功能将自动检测出失败的节点并根据需要执行恢复操作。随着自动恢复，我们能够从失败中平稳恢复，并在没有任何人为干预的情况下继续训练，避免昂贵的停机时间和手动管理。开始后，无烦恼！

软件效率

软件配置优化一直是一个大问题。幸运的是，我们基于Pytorch的作曲家库可以最大程度地提高培训效率。与上一轮实验相似，随着GPU的数量的增加，作曲家继续保持出色的吞吐量缩放能力。在此更新中，我们添加了进一步的优化（低精度的集体态和低精度分层，并行完全碎片的数据），以实现接近完美的功能缩放功能，将工作扩展到最高128 GPU，从而将成本保持在50,000美元以内。我们还使用作曲家的本机指数移动平均线（EMA）算法来启用培训结束（在第二阶段的800,000次迭代中）的EMA，从而节省了相当大的内存和训练计算能力。

管理 100 TB 数据

我们在培训中使用了一个含有7.9亿个样品的LAION-5B，总数据量超过100 TB。巨大的数据集尺寸使得很难管理，尤其是在有必要与具有独立本地存储的多群落合作时。

Mosaicml StreamDataSet库使处理大量数据集更加容易，更快。图书馆提供的三个核心功能在此培训中也起着关键作用：

混合存储在不同位置的数据集。我们根据图像分辨率将每个样本存储在不同的数据集中。在培训期间，我们使用Mosaicml流媒体库来混合每个数据集的分辨率材料。

2。在即时轮子中恢复。我们可以在一轮中立即恢复培训，这相当于在训练过程中意识到整个数据集的“断点连续传输”，从而大大节省了整个时间。

3。以灵活的方式确定性。 MosaiCML StreamDataSet库能够以确定性的方式进行调整数据，而不会受到训练GPU的影响。这使我们能够准确地重现训练效果并大大简化调试步骤。

人类评估结果

通常很难评估图像生成模型的实际性能，除了投资人力之外，别无他法。在盲试评估中，我们测量了用户对图像质量的偏好，并在稳定扩散2和我们自己的扩散模型之间进行了迅速的单词对齐。根据用户偏好，我们得出结论，两组模型质量相等（见图5）。所有图像均根据Drawbench基准测试生成，揭示了成像纸中提出的单词。

图5：我们对图像质量的手动评估结果（左）和揭示单词对齐（右）的结果。误差线显示的置信区间为95。在这两个实验中，两组模型之间用户偏好的差异与测量的不确定性相当，因此我们得出结论，两组模型的总体质量也基本相同。

未来展望

写在最后

以下是我们扩散模型产生的图像结果。团队成员都很开心，希望您能在其中找到自己的乐趣。

原始链接：

https://www.mosaicml.com/blog/training-stable-diffusion-from-scratch-part-2