Jupyterhub配置限制张量流gpu内存

Question

我正在为我的学生在课堂上建立一个与Jupyterhub（docker spawner）的张量流环境，但我遇到了这个问题。

  

默认情况下，TensorFlow几乎映射了所有GPU的所有GPU内存   （受CUDA_VISIBLE_DEVICES约束）对过程可见。这个完成了   更有效地使用相对宝贵的GPU内存资源   通过减少内存碎片在设备上。   （来自https://www.tensorflow.org/tutorials/using_gpu）

如果班上有人使用python程序和gpu，那么gpu内存几乎耗尽。根据这种情况，我需要手动添加一些限制代码。 像：

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

但这不是一个很好的解决方案。我应该在每次生成新代码时添加此代码。

如果jupyterhub可以添加一些配置来避免这种情况或其他很好的解决方案？ 请让我知道，谢谢！

Answer 1

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

这很好用