如何防止张量流分配GPU内存的总量?
python
tensorflow
5
0

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

对于中小型型号,Titan X的12 GB通常足以让2-3个人在同一GPU上同时进行训练。如果模型足够小,以至于单个模型无法充分利用GPU的所有计算单元,那么与运行另一个训练过程相比,这实际上可以提高速度。即使在同时访问GPU确实减慢了单个训练时间的情况下,具有让多个用户同时在GPU上进行训练的灵活性仍然很好。

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

如果有人知道对于一个给定的模型就足够了,有没有办法使TensorFlow仅分配4 GB的GPU内存?

参考资料:
Stack Overflow
收藏
评论
共 10 个回答
高赞 时间 活跃

以上所有答案均假设通过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)

同样在“急切”模式(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)

编辑:2018年11月4日,例如,如果您要使用tf.contrib.gan.train ,则可以使用类似于波纹管的东西:

tf.contrib.gan.gan_train(........, config=conf)
收藏
评论

这是Deep Learning with TensorFlow书摘录

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

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

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

第二种方法是per_process_gpu_memory_fraction选项,该选项确定应为each可见GPU分配的总内存量的一部分。 注意:不需要释放内存,完成后甚至会加剧内存碎片。

2)分配固定内存

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

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

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

收藏
评论
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

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

收藏
评论

您可以通过在可选config参数中传递tf.GPUOptions来设置构造tf.Session时要分配的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进行设置。

收藏
评论

对于Tensorflow 2.0版,请使用以下代码段

 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)
收藏
评论

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,请参阅: 此答案

收藏
评论

已针对TensorFlow 2.0 Alpha及更高版本进行了更新

从2.0 Alpha文档来看,答案现在只有一行,然后您可以使用TensorFlow进行任何操作:

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)
收藏
评论

您可以使用

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();
}
收藏
评论

对于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)
收藏
评论

无耻插件:如果安装了GPU支持的Tensorflow,则无论您将其设置为仅使用CPU还是GPU,该会话都会首先分配所有GPU。我可能会补充一点,即使您将图形设置为仅使用CPU,也应该设置相同的配置(如上回答:)),以防止不必要的GPU占用。

并且在类似IPython的交互式界面中,您还应该设置configure,否则它将分配所有内存,而几乎不分配给其他人。有时很难注意到。

收藏
评论
新手导航
  • 社区规范
  • 提出问题
  • 进行投票
  • 个人资料
  • 优化问题
  • 回答问题

关于我们

常见问题

内容许可

联系我们

@2020 AskGo
京ICP备20001863号