文章目录
原创声明
本文为 HinGwenWoong 原创,如果这篇文章对您有帮助,欢迎转载,转载请阅读文末的【授权须知】,感谢您对 HinGwenWoong 文章的认可!
前言
部署深度学习服务的时候,往往不是让其吃满一整张卡,而且有时候会出现致命的 OOM
(Out of Memory)错误,这就需要适当限制下显存,下面介绍下如何使用代码限制显存。
关于在TF2.0中使用GPU的详细说明可以从 官方文档 中找到。在本文中,我将向您展示一些您可以立即使用的文档中的代码。
我是 HinGwenWoong,一个有着清晰目标不停奋斗的程序猿,热爱技术,喜欢分享,码字不易,如果帮到您,请帮我在屏幕下方点赞 👍,您的点赞可以让技术传播得更远更广,谢谢!
一、TensorFlow 2.x
方法一: set_memory_growth
该方法是设置 显存按需增长 来实现一开始不吃满来达到目的。
import tensorflow as tf
gpu_list = tf.config.experimental.list_physical_devices('GPU')
if len(gpu_list) > 0 :
for gpu in gpu_list:
try:
# 设置多张 GPU ,如果不需要 for 去设置多张,则使用 list 的索引设置即可
tf.config.experimental.set_memory_growth(gpu, True)
except RuntimeError as e:
print(e)
else:
print("Got no GPUs")
注意:
- 所有 GPU 的【显存按需增长】的配置应该相同。
- 应该在初始化 GPU 之前设置【显存按需增长】。
方法二:memory_limit
该步骤将第一个 GPU 的内存使用量限制为 2048MB
。只需根据需要更改 gpu_list
的索引和 memory_limit
的设置即可。
import tensorflow as tf
using_gpu_index = 0 # 使用的 GPU 号码
gpu_list = tf.config.experimental.list_physical_devices('GPU')
if len(gpu_list) > 0:
try:
tf.config.experimental.set_virtual_device_configuration(
gpu_list[using_gpu_index],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)]
)
except RuntimeError as e:
print(e)
else:
print("Got no GPUs")
二、TensorFlow 1.x
方法一:allow_growth
下面这段代码对应于 TF2.x
的【方法一】。
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
方法二:per_process_gpu_memory_fraction
下面这段代码对应 TF2.x
的【方法二】,但它设置的是显存的百分比,不是一个确定的值。下面的例子设置了使用60%
的显存
import tensorflow as tf
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.6) # 改变这个百分比即可
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
总结
以上就是 TF 2.x 和 1.x 显存的限制手段,希望能够帮助大家部署更加平稳。
我是 HinGwenWoong,一个有着清晰目标不停奋斗的程序猿,热爱技术,喜欢分享,码字不易,如果帮到您,请帮我在屏幕下方点赞 👍,您的点赞可以让技术传播得更远更广,谢谢!
授权须知
- 原创文章在推送一天后才可进行转载
- 转载文章,禁止声明原创
- 不允许直接二次转载,转载请根据原文链接联系作者
- 若无需改版,在文首清楚标注作者及来源/原文链接,并删除【原创声明】,即可直接转载。
但对于未注明转载来源/原文链接的文章,我将保留追述的权利。作者:HinGwenWoong
一个有着清晰目标不停奋斗的程序猿,热爱技术,喜欢分享,共同进步!
CSDN: HinGwenWoong
原文链接:深度学习 | TensorFlow 2.x 和 1.x 限制显存(超详细) - 若需要修改文章的排版,请根据原文链接联系作者
- 再次感谢您的认可,转载请遵守如上转载须知!