如何限制TF Slim中的GPU内存使用？

Question

当使用TF Slim的train_image_classifier.py进行训练时，我想告诉Slim仅分配所需的GPU内存，而不分配所有内存。

如果我使用TF而不是Slim，我可以这样说：

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

甚至仅仅是为了对GPU内存的使用设置硬性限制：

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

我该如何告诉Slim同样的事情？

我的理解失败是Slim似乎使用了它自己的循环，而我在配置循环的细节上找不到文档。 因此，即使有人可以将我指向优秀的Slim文档，也将是很棒的。

提前致谢！

Answer 1

您可以通过session_config参数传递allow_growth选项，该参数传递给train方法，如下所示：

session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth = True
slim.learning.train(..., session_config=session_config)

参见tensorflow / contrib / slim / python / slim / learning.py＃L615和tensorflow＃5530 。