如何在我的代码之外设置 tensorflow gpu 内存分配？答案

【问题标题】：How to set tensorflow gpu memory allocation outside of my code?如何在我的代码之外设置 tensorflow gpu 内存分配？
【发布时间】：2018-04-28 08:04:50
【问题描述】：

我在 spark 上使用 tensorflow。我的 tensorflow 分配了几乎所有的 GPU 内存（10GB 或 12GB）。只有 'import tensorflow' 占用所有内存我知道如何在代码中配置 gpu 内存分数。我知道的方法是添加下面提到的代码

import tensorflow as tf
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.2)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

但是..在火花上，即使我在代码中添加了这个配置，它也不能应用于其他 GPU 机器.. 所以，我想知道如何在代码之外配置或设置 gpu 分数，这意味着通过 tensorflow 本身..

【问题讨论】：

标签： apache-spark memory tensorflow gpu

【解决方案1】：

在tensorflow/tensorflow/core/common_runtime/gpu/gpu_device.cc 变化自

if (config_memory_fraction == 0) {
allocated_memory = available_memory;

到

if (config_memory_fraction == 0) {
allocated_memory = total_memory * 0.2(fraction u want to use);

然后，通过 bazel 构建它。

这有点硬编码。如果你修改源代码，你可以做gpu内存的动态分配

【讨论】：