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

Question

我试图控制为一个张量流量估算器tf.estimator.Estimator分配的GPU内存的大小。目的是仅在同一GPU上分配一半来运行其他tensorflow网络。我找到了贡献版本但不是官方版本。有人知道它是否可能？

Answer 1

创建Estimator实例时，您可以传入构造函数config和tf.estimator.RunConfig实例。 RunConfig有一个session_config属性，可用于设置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)