search

ZZZHANG-jx / DocRes

[CVPR 2024] DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks
MIT License
287 stars 26 forks source link

end2end 显存递增 #1

Closed yazheng0307 closed 4 months ago

yazheng0307 commented 4 months ago

感谢作者的工作！ 我在inference.py脚本中使用end2end的任务时，发现显存在model_init()后占用一部分（约1.9G），再进行到deshadowing任务时显存的占用明显增加，直到OOM，尝试使用torch.cuda.empty_cache()未果，请教一下解决方案，再次感谢作者团队的工作！

xuhangc commented 4 months ago

应该是因为你的图比MAX_SIZE要大，建议resize一下

ZZZHANG-jx commented 4 months ago

我们是用48G显存的A6000推理的，如果显卡没那么大显存，可以把以下几行的MAX_SIZE调小一点 https://github.com/ZZZHANG-jx/DocRes/blob/d09e6d7b1e12f02e19fe1223c5208a222fe4a507/inference.py#L131 https://github.com/ZZZHANG-jx/DocRes/blob/d09e6d7b1e12f02e19fe1223c5208a222fe4a507/inference.py#L170

yazheng0307 commented 4 months ago

好的！谢谢