如果你正在使用Tensorflow 2，请尝试以下步骤:

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

当你构造一个tf时，你可以设置GPU内存的分配比例。会话通过传递一个tf。GPUOptions作为可选配置参数的一部分:

# 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基础上设置这个。

以上答案都是指在TensorFlow 1中设置一定的内存。或者在TensorFlow 2.X中允许内存增长。

方法tf.config.experimental。Set_memory_growth确实适用于在分配/预处理期间允许动态增长。然而，人们可能喜欢从一开始就分配一个特定的GPU内存上限。

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

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

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)]

我尝试在voc数据集上训练unet，但由于图像大小巨大，内存结束。我尝试了上面所有的技巧，甚至尝试了batch size==1，但没有任何改善。有时候TensorFlow版本也会导致内存问题。尝试使用

PIP install tensorflow-gpu==1.8.0

对于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)

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

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