【stable-diffusion企业级教程04】EMA你走,拥抱16G显存!Xformers...
发布网友
发布时间:2024-10-16 06:19
我来回答
共1个回答
热心网友
时间:2024-10-16 06:24
回顾了上一讲中成功使用deepspeed框架,将sd模型以fp16的精度训练,显存消耗降低至18G左右的情况。然而,有读者提出自己仅有16G显卡,希望进一步降低训练时的资源需求。今天,我们将介绍两个方法,旨在进一步降低整体显存需求。
首先,我们来了解指数移动平均(EMA)。熟悉股票技术分析的读者对移动平均(MA)应不陌生,它通过计算前K天的值的平均值作为指标。在深度学习领域中,EMA是一种更新模型权重的方法,它通过维护一个影子权重,对模型参数进行“加权平均”,以在最终评估集上获得更好的效果。个人理解,batch_gradient_decent是不同样本共同决定更新方向;而EMA则是跨batch决定更新幅度。
接下来,我们讨论模型选择。通过链接提供的两个模型,我们发现加载包含EMA权重的模型(如sd-v1-4-full-ema.ckpt)会导致额外参数的保存,这增加了加载模型时的显存消耗。为了解决这个问题,我们选择不加载EMA权重的基线模型(如sd-v1-4.ckpt)。
在训练设置中,我们避免生成与EMA相关的权重。通过观察代码,我们发现EMA权重通过on_train_batch_end方法实现。我们只需在配置文件的对应位置设置成员变量self.use_ema为False,从而在训练时避免生成EMA权重。
通过上述更改,我们在两张卡上重新启动训练,选择batchsize=2。结果显示,成功将显存消耗降低至16G以下。
接下来,我们引入xformers库,以进一步加速训练过程。xformers是Facebook开源的计算库,能够更有效地实现transformer中的计算。在2.X模型中,xformers被引入以提高性能。通过github项目搜索,我们了解到xformers主要在ldm中的相关文件和配置文件中使用。在引入xformers后,我们观察到在特定条件下(如较高分辨率图像)显存优化效果更为显著。使用xformers时,需确保推理过程也运行在包含xformers的环境中,以避免精度溢出导致的黑图问题。
综上所述,我们通过移除EMA参数和配置,将训练时batchsize=2的显存消耗降至16G以下,并通过引入xformers为更高分辨率图像处理做准备。代码已更新至仓库,欢迎参考使用。