343

我在一个计算资源共享的环境中工作,也就是说,我们有几台服务器机器,每台机器都配备了一些 Nvidia Titan X GPU。

对于中小型模型,Titan X 的 12 GB 通常足以让 2-3 人在同一个 GPU 上同时运行训练。如果模型足够小,以至于单个模型不能充分利用 GPU 的所有计算单元,那么与一个接一个地运行训练过程相比,这实际上会导致加速。即使在并发访问 GPU 确实会减慢个人训练时间的情况下,让多个用户同时在 GPU 上训练的灵活性仍然很好。

TensorFlow 的问题在于,默认情况下,它会在启动时分配全部可用的 GPU 内存。即使对于一个小型的两层神经网络,我看到所有 12 GB 的 GPU 内存都用完了。

有没有办法让 TensorFlow 只分配 4 GB 的 GPU 内存,如果知道这对于给定模型来说已经足够了吗?

4

16 回答 16

326

tf.Session您可以通过将 atf.GPUOptions作为可选config参数的一部分传递来设置在构造 a 时要分配的 GPU 内存的比例:

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

per_process_gpu_memory_fraction作为 GPU 内存量的硬上限,该内存量将由同一台机器上的每个 GPU 上的进程使用。目前,这一比例统一应用于同一台机器上的所有 GPU;无法在每个 GPU 的基础上进行设置。

于 2015-12-10T11:00:19.170 回答
210
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

https://github.com/tensorflow/tensorflow/issues/1578

于 2016-05-26T07:43:45.363 回答
83

对于 TensorFlow 2.0 和 2.1文档):

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)

对于 TensorFlow 2.2+文档):

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

文档还列出了更多方法:

  • 将环境变量设置TF_FORCE_GPU_ALLOW_GROWTHtrue.
  • 用于tf.config.experimental.set_virtual_device_configuration在虚拟 GPU 设备上设置硬限制。
于 2019-04-05T18:26:37.670 回答
58

这是本书的摘录Deep Learning with TensorFlow

在某些情况下,希望进程仅分配可用内存的子集,或者仅在进程需要时增加内存使用量。TensorFlow在会话上提供了两个配置选项来控制这一点。第一个是allow_growth选项,它尝试根据运行时分配只分配尽可能多的 GPU 内存,它开始分配非常少的内存,随着会话开始运行并且需要更多的 GPU 内存,我们扩展了 TensorFlow 所需的 GPU 内存区域过程。

1)允许增长:(更灵活)

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

第二种方法是选项,它决定了应该分配可见 GPUper_process_gpu_memory_fraction的内存总量的比例。注意:不需要释放内存,完成后甚至会加剧内存碎片。each

2)分配固定内存

仅通过以下方式分配40%每个 GPU 的总内存:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

注意: 这仅在您确实想要绑定 TensorFlow 进程上可用的 GPU 内存量时才有用。

于 2018-01-11T18:57:16.453 回答
26

对于 TensorFlow 2.0 和 2.1 版,使用以下代码段

 import tensorflow as tf
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_memory_growth(gpu_devices[0], True)

对于以前的版本,以下片段曾经对我有用:

import tensorflow as tf
tf_config=tf.ConfigProto()
tf_config.gpu_options.allow_growth=True
sess = tf.Session(config=tf_config)
于 2019-12-01T14:47:55.657 回答
22

上面的所有答案都假设通过sess.run()调用执行,这正在成为 TensorFlow 最新版本中的例外而不是规则。

使用tf.Estimator框架(TensorFlow 1.4 及更高版本)时,将分数传递给隐式创建MonitoredTrainingSession的方法是,

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
trainingConfig = tf.estimator.RunConfig(session_config=conf, ...)
tf.estimator.Estimator(model_fn=..., 
                       config=trainingConfig)

同样在 Eager 模式下(TensorFlow 1.5 及更高版本),

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
tfe.enable_eager_execution(config=conf)

编辑:11-04-2018 例如,如果您要使用tf.contrib.gan.train,那么您可以使用类似于下面的内容:

tf.contrib.gan.gan_train(........, config=conf)
于 2018-02-08T03:25:04.283 回答
13

您可以使用

TF_FORCE_GPU_ALLOW_GROWTH=true

在您的环境变量中。

张量流代码中:

bool GPUBFCAllocator::GetAllowGrowthValue(const GPUOptions& gpu_options) {
  const char* force_allow_growth_string =
      std::getenv("TF_FORCE_GPU_ALLOW_GROWTH");
  if (force_allow_growth_string == nullptr) {
    return gpu_options.allow_growth();
}
于 2019-06-02T17:15:29.323 回答
12

Tensorflow 2.0 Beta 和(可能)更高版本

API 再次更改。现在可以在以下位置找到:

tf.config.experimental.set_memory_growth(
    device,
    enable
)

别名:

  • tf.compat.v1.config.experimental.set_memory_growth
  • tf.compat.v2.config.experimental.set_memory_growth

参考:

另请参阅: Tensorflow - 使用 GPUhttps ://www.tensorflow.org/guide/gpu

对于 Tensorflow 2.0 Alpha,请参阅: this answer

于 2019-06-17T13:08:14.587 回答
8

上面的所有答案都是指在版本中将内存设置到一定程度TensorFlow 1.X或允许内存增长TensorFlow 2.X

该方法 tf.config.experimental.set_memory_growth确实适用于在分配/预处理期间允许动态增长。尽管如此,人们可能希望从一开始就分配一个特定上限的 GPU 内存。

分配特定 GPU 内存背后的逻辑也是为了防止在训练期间出现 OOM 内存。例如,如果一个人在打开显存消耗 Chrome 选项卡/任何其他视频消耗进程时进行训练,则tf.config.experimental.set_memory_growth(gpu, True)可能会导致引发 OOM 错误,因此在某些情况下需要从一开始就分配更多内存。

在 TensorFlow 2.X 中为每个 GPU 分配内存的推荐且正确的方法是通过以下方式完成的:

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]
于 2020-08-28T09:54:59.383 回答
7

如果您使用的是 Tensorflow 2,请尝试以下操作:

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.Session(config=config)
于 2020-03-12T06:44:53.993 回答
7

无耻插件:如果你安装了GPU支持的Tensorflow,无论你设置为只使用CPU还是GPU,session都会先分配所有GPU。我可能会添加我的提示,即使您将图表设置为仅使用 CPU,您也应该设置相同的配置(如上面回答的:))以防止不必要的 GPU 占用。

在像 IPython 和 Jupyter 这样的交互式界面中,您还应该设置该配置,否则,它将分配所有内存而几乎没有留给其他人。这有时很难注意到。

于 2017-05-23T07:52:55.927 回答
5

对于Tensorflow 2.0,这个解决方案对我有用。(TF-GPU 2.0、Windows 10、GeForce RTX 2070)

physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
tf.config.experimental.set_memory_growth(physical_devices[0], True)
于 2019-10-05T21:05:50.780 回答
3
# allocate 60% of GPU memory 
from keras.backend.tensorflow_backend import set_session
import tensorflow as tf 
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6
set_session(tf.Session(config=config))
于 2019-09-27T00:28:35.503 回答
3

这段代码对我有用:

import tensorflow as tf
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.compat.v1.InteractiveSession(config=config)
于 2020-07-08T22:30:04.190 回答
2

好吧,我是 tensorflow 的新手,我有 Geforce 740m 或带有 2GB ram 的 GPU,我正在运行 mnist 手写类型的母语示例,训练数据包含 38700 个图像和 4300 个测试图像,并试图获得精度、召回率、 F1 使用以下代码作为 sklearn 并没有给我精确的结果。一旦我将它添加到我现有的代码中,我就开始出现 GPU 错误。

TP = tf.count_nonzero(predicted * actual)
TN = tf.count_nonzero((predicted - 1) * (actual - 1))
FP = tf.count_nonzero(predicted * (actual - 1))
FN = tf.count_nonzero((predicted - 1) * actual)

prec = TP / (TP + FP)
recall = TP / (TP + FN)
f1 = 2 * prec * recall / (prec + recall)

再加上我猜我的模型很重,我在 147、148 个 epoch 之后出现内存错误,然后我想为什么不为任务创建函数所以我不知道它是否在 tensrorflow 中以这种方式工作,但我想如果一个局部变量是使用过,当超出范围时它可能会释放内存,并且我在模块中定义了上述用于训练和测试的元素,我能够毫无问题地实现 10000 个 epoch,我希望这会有所帮助..

于 2019-01-21T17:26:59.050 回答
1

我试图在 voc 数据集上训练 unet,但由于图像尺寸巨大,内存完成。我尝试了上述所有技巧,甚至尝试了批量大小==1,但没有任何改善。有时 TensorFlow 版本也会导致内存问题。尝试使用

pip install tensorflow-gpu==1.8.0

于 2018-10-16T06:05:52.857 回答