config.py 文件源码

def get_session_config():
  import tensorflow as tf
  session_args = {
      'intra_op_parallelism_threads': CONFIG['nthread'],
      'inter_op_parallelism_threads': CONFIG['ncpu'],
      'allow_soft_placement': True,
      'log_device_placement': CONFIG['debug'],
  }
  if CONFIG['ngpu'] > 0:
    if CONFIG['cnmem'] > 0:
      session_args['gpu_options'] = tf.GPUOptions(
          per_process_gpu_memory_fraction=CONFIG['cnmem'],
          allow_growth=False)
    else:
      session_args['gpu_options'] = tf.GPUOptions(
          allow_growth=True)
  return session_args