bytedance / bytedance/StyleSSP
Tips for reducing video memory overhead to below 24GB
Open
- Dominant language
- Python
- Stars
- 92
- Forks
- 6
- PR merge metrics
- No merged PRs in 30d
Description
24GB显存对于小实验组是一个难以突破的阈值(如只有3090或4090的组),但当前的infer_style.py代码产生的显存开销很容易突破这个阈值,导致无法运行该代码。
这里提一个很简易但能有效将显存开销降至20G以下的改动。
对于 model 可以在执行 generate_caption函数后通过 model.cpu()的方式手动卸载入cpu,在下次运行前,通过model.cuda()的方法手动载入GPU;
对于所需使用的几个pipeline(行171,行256,行386所创建的几个pipeline),可以通过以下代码将使pipeline所使用的组件自动装载入GPU和卸载入CPU,从而降低显存开销。
`pipe.enable_model_cpu_offload()`
`pipe.to(“cpu”)`
完成上述处理后,可以将 行276 的代码删除掉以便于在循环运行生成代码时无需重新载入所用的几个模型。
`del pipe_inversion, pipe_inference_, model`
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.