python - 防止为 MonitoredTrainingSession 分配 GPU 内存

Question

我正在尝试限制 MonitoredTrainingSession 中的 GPU 内存分配。

tf.GPUOptions的设置方法如下:How to prevent tensorflow from allocating the totality of a GPU memory?在 MonitoredTrainingSession 的情况下不起作用。

我尝试过:

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=.1)
# or allow_growth=True
config = tf.ConfigProto(allow_soft_placement=False,
                        device_filters=filters,
                        gpu_options=gpu_options)

scaffold = tf.train.Scaffold(saver=tf.train.Saver(max_to_keep=100, keep_checkpoint_every_n_hours=.5))

with tf.train.MonitoredTrainingSession(
                server.target,
                is_chief=True,
                checkpoint_dir=log_dir,
                scaffold=scaffold,
                save_checkpoint_secs=600,
                save_summaries_secs=30,
                log_step_count_steps=int(1e7),
                config=config) as session:

尽管使用 tf.GPUOptions 内存消耗为 10189MiB/11175MiB

Answer 1

我找出了问题所在:打开的第一个 session 需要包含内存选项。

因此，如果有疑问，请在脚本开头打开一个具有内存限制的虚拟 session 。