python - 如何在 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 文档，那也很棒。

提前致谢！

score 1 · Accepted Answer

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

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

1 回答 1