设置GPU显存的方法和我们通常用tf.Session(config=tf_config)一样,可以通过tf_config.gpu_options.allow_growth = True设置为自动分配显存,或者用tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5指定百分比。
区别在于,调用时这里不能直接指定session_config,需要用replace的方法把config设置到run_config中。
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
#tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 #如果是手动指定百分比
run_config = tf.contrib.learn.RunConfig(
model_dir=Config.train.model_dir,
save_checkpoints_steps=Config.train.save_checkpoints_steps,
).replace(session_config=tf_config)
tf.contrib.learn.learn_runner.run(
experiment_fn=experiment_fn,
run_config=run_config,
schedule=mode,
hparams=params
)