如何限制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

您可以通过传递给allow_growth方法的session_config参数传递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。