本文为 HinGwenWoong 原创，如果这篇文章对您有帮助，欢迎转载，转载请阅读文末的【授权须知】，感谢您对 HinGwenWoong 文章的认可！
部署深度学习服务的时候，往往不是让其吃满一整张卡，而且有时候会出现致命的 OOM （Out of Memory）错误，这就需要适当限制下显存，下面介绍下如何使用代码限制显存。
我是 HinGwenWoong，一个有着清晰目标不停奋斗的程序猿，热爱技术，喜欢分享，码字不易，如果帮到您，请帮我在屏幕下方点赞 👍，您的点赞可以让技术传播得更远更广，谢谢！
方法一： 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:
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_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")
方法一： 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，一个有着清晰目标不停奋斗的程序猿，热爱技术，喜欢分享，码字不易，如果帮到您，请帮我在屏幕下方点赞 👍，您的点赞可以让技术传播得更远更广，谢谢！
Original: https://blog.csdn.net/hxj0323/article/details/121629067
Author: HinGwenWoong
Title: 深度学习 | TensorFlow 2.x 和 1.x 限制显存（超详细）
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/497077/
转载文章受原作者版权保护。转载请注明原作者出处！