I'm trying to control the size of GPU memory allocated for one tensorflow estimator tf.estimator.Estimator. The purpose is to only allocate half to run other tensorflow net on the same GPU. I found for the contrib version but not for the official. Someone knows if it's possible?
1 Answers
When you create an Estimator instance, you can pass in the constructor's config a tf.estimator.RunConfig instance.
The RunConfig has a session_config attribute you can use to set a tf.ConfigProto with the session's parameters.
In code, this translates to:
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)
Related