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

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

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

How to limit GPU memory use in TF Slim?