【用时间换显存】RG版DFL（解压即用）

lispmox · 发表于 2021-9-8 06:14:58

更新：如果显存有限，建议把 place models and optimizer on GPU 设为no。训练时模型权重梯度优化器状态也需要消耗存储空间，大概是模型权重的3-4倍。比如参数为 “256 wf liae-ud 256 64 64 22”的模型权重就有554M，优化器状态有1.08G，加上梯度以及元信息耗费的空间，大概有3G左右的显存用于存模型不能用于计算（不一定准确，3G是用pytorch加载等价模型后看nvidia-smi得到的。tensorflow是静态图调度可能更省一点，但应该差不了多少）。如果显存比较大，这点占用可以忽略不记，但小显存就省着点用。

更新：调整一下价格，让更多人能使用。话说灵石有办法退吗，本贴有两位提前上车了。

前段时间我发了一个在6G显存上跑512 DF-UD模型的帖子（链接），主要思路就是将DFL模型导入pytorch框架，然后借助checkpoint和混合精度来减少显存占用。实测下来可以用bs8跑默认参数的512分辨率DF-UD模型。一开始我是抱着学习原理的心态去基于pytorch重写DFL的部分代码，后来有人希望我能制作一个集成软件，也有人私信我想要代码。由于个人精力有限，我不想花费大量时间来重新开发并维护一个“换脸”软件，哪怕只是其中一个子模块，所以抱歉了各位。

话虽如此，为了响应部分同志的需求，我决定在DFL原版软件的基础上进行适当修改，使其能够尽量减少训练期间的显存占用。代价就是牺牲训练的速度，因此比较适合用来复用预训练模型（丹）。实测下来可以用bs4跑默认参数的512分辨率DF-UD模型，比之前基于pytorch的实现指标差一些，但可以使用DFL完整的功能。

其实用中低端显卡跑太大的模型挺勉强的，速度很慢。这个版本的好处可能就是让低显存用户也能开高一些参数吧，或者256分辨率下开启GAN。

>>> 从论坛里找到预训练模型 512 DF-UD

>>> 滚石大大的万能模型
QQ截图20210907195044.jpg

PS：GAN模型在8月份有改动，老模型好像不能复用了

改动：
1. DFL版本应该是0803，直接从github上拉取的最新镜像。官方集成的版本相较于github版本有一些修改，因此细节上略有不同。
2. 基于conda环境打包tensorflow2.5.0+cuda11.3.1+cudnn8.2.1，在我的电脑上（win10 64位 20系列显卡）比官方3000 series更稳定一点。
3. ffmpeg更新到4.4版本，修复了一些问题。官方集成版本合成视频在win10自带播放器上可能会黑屏。
4. 修改了环境变量脚本，给后面也想打包python的同学一点参考。
5. 分享里附带一个recompute.ipynb文件，包含一个在tensorflow2+gpu+graph mode的环境下正常工作的recompute_grad，不能用在默认的eager mode下。

关于recompute.ipynb，里面包含recompute_grad函数的实现，用于在静态计算图上生成检查点。我期望在tensorflow2.x的版本上实现gradient checkpoint功能，以此减少显存占用。但是网上能找到的以及默认的tf.recompute_grad只能在eager mode下工作，而DFL工作在graph mode下。openai之前开源过一个graph mode版本（链接），可惜只能运行在tensorflow1.x上。既然tensorflow2.x是未来，同时也为了使用cuda11并支持30系列显卡，我尝试自己实现了一个recompute_grad函数，能工作在tensorflow2.x的graph mode下。代码很简洁，我不保证是最优实现，但确实可以节省一定的显存占用。因为个人对tensorflow框架并不熟悉，期待更好的解决方案。

购买主题 已有 121 人购买 本主题需向作者支付 2 灵石 才能浏览