Reputation:
When training using TF Slim's train_image_classifier.py I would like to tell Slim to only allocate what GPU memory it needs, rather than allocating all the memory.
Were I using straight up TF and not Slim I could say this:
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
Or even just this to put a hard cap on GPU memory use:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
How can I tell Slim the same thing(s)?
My comprehension fail is that Slim seems to use it's own loop and I can't find docs on the nitty gritty of configuring the loop. So, even if someone could point me to good Slim docs that'd be fantastic.
Thanks in advance!
Upvotes: 0
Views: 209
Reputation: 1703
You can pass the allow_growth
option via the session_config
parameter that is passed to the train
method as follow:
session_config = tf.ConfigProto()
session_config.gpu_options.allow_growth = True
slim.learning.train(..., session_config=session_config)
See tensorflow/contrib/slim/python/slim/learning.py#L615 and tensorflow #5530.
Upvotes: 0