深度学习 | TensorFlow 2.x 和 1.x 限制显存（超详细）

---

原创声明

本文为 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，一个有着清晰目标不停奋斗的程序猿，热爱技术，喜欢分享，码字不易，如果帮到您，请帮我在屏幕下方点赞 👍，您的点赞可以让技术传播得更远更广，谢谢！

---

更多阅读推荐

• 深度学习实战 | 智慧工地安全帽和危险区域检测系统（代码已开源！）
• 深度学习OCR | 文字识别网络 CRNN 论文翻译

---

授权须知

1. 原创文章在推送一天后才可进行转载
2. 转载文章，禁止声明原创
3. 不允许直接二次转载，转载请根据原文链接联系作者
4. 若无需改版，在文首清楚标注作者及来源/原文链接，并删除【原创声明】，即可直接转载。
   但对于未注明转载来源/原文链接的文章，我将保留追述的权利。

   作者：HinGwenWoong
   一个有着清晰目标不停奋斗的程序猿，热爱技术，喜欢分享，共同进步！
   CSDN: HinGwenWoong
   原文链接：深度学习 | TensorFlow 2.x 和 1.x 限制显存（超详细）

5. 若需要修改文章的排版，请根据原文链接联系作者
6. 再次感谢您的认可，转载请遵守如上转载须知！