设置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中。
1 tf_config = tf.ConfigProto()
2 tf_config.gpu_options.allow_growth = True
3 #tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 #如果是手动指定百分比
4 run_config = tf.contrib.learn.RunConfig(
5 model_dir=Config.train.model_dir,
6 save_checkpoints_steps=Config.train.save_checkpoints_steps,
7 ).replace(session_config=tf_config)
8
9 tf.contrib.learn.learn_runner.run(
10 experiment_fn=experiment_fn,
11 run_config=run_config,
12 schedule=mode,
13 hparams=params
14 )