用于限制 tensorflow gpu 内存的 Jupyterhub 配置

Question

我正在用 Jupyterhub(docker spawner) 为 class 的学生构建一个 tensorflow 环境，但我遇到了这个问题。

By default, TensorFlow maps nearly all of the GPU memory of all GPUs (subject to CUDA_VISIBLE_DEVICES) visible to the process. This is done to more efficiently use the relatively precious GPU memory resources on the devices by reducing memory fragmentation. (from https://www.tensorflow.org/tutorials/using_gpu)

如果class有人用gpu使用python程序，那么gpu内存将接近exhaust.According这种情况，我需要手动添加一些限制代码。喜欢：

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

但这不是很好solution.I每次生成新代码时都应该添加此代码。

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

这个效果很好

用于限制 tensorflow gpu 内存的 Jupyterhub 配置

Jupyterhub config for limit tensorflow gpu memory

gpu

jupyter

tensorflow

jupyterhub