Как ограничить использование памяти GPU в TF Slim?

Question

При обучении с использованием TF Slim, train_image_classifier.py, я хотел бы сказать Слиму, чтобы он выделял только ту память, которая ему нужна, а не всю память.

Если бы я использовал прямой TF, а не Slim, я мог быскажем так:

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

Или даже просто, чтобы жестко ограничить использование памяти графическим процессором:

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

Как я могу сказать Слиму то же самое?

Мой сбой в понимании заключается в том, что Слим, похоже, использует свой собственный цикл, и я не могу найти документы по мелочам при настройке цикла.Так что, даже если бы кто-то мог указать мне на хорошие Slim документы, это было бы замечательно.

Заранее спасибо!

Answer 1

Вы можете передать параметр allow_growth с помощью параметра session_config, который передается методу train, следующим образом:

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

См. tenorflow / contrib / slim / python / slim / learning.py # L615 и тензор потока # 5530 .