Jupyterhub config for limit tensorflow gpu memory

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

这很好用