tensorflow - 如何使用 tf.estimator 控制 GPU 内存大小

我正在尝试控制为一个 tensorflow estimator tf.estimator.Estimator 分配的 GPU 内存的大小。目的是只分配一半在同一 GPU 上运行其他 tensorflow 网络。我找到了 contrib 版本，但没有找到官方版本。有人知道这是否可能？

最佳答案

创建 Estimator 实例时，可以在构造函数的 config 中传入一个 tf.estimator.RunConfig 实例。RunConfig 有一个 session_config 属性，您可以使用它来设置带有 session 参数的 tf.ConfigProto。

在代码中，这转化为:

session_config = tf.ConfigProto()
session_config.gpu_options.per_process_gpu_memory_fraction = 0.5
estimator_config = tf.estimator.RunConfig(session_config=session_config)
my_estimator = tf.estimator.Estimator(..., config=estimator_config)

关于tensorflow - 如何使用 tf.estimator 控制 GPU 内存大小，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/47389988/