Skip to content

GitLab

Commit e66575c7 authored by Haoyu Zhang's avatar Haoyu Zhang
Browse files

Changing default GPU private thread count to 1

parent eebea3f8
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
official/resnet/keras/keras_common.py
View file @ e66575c7
...@@ -115,7 +115,7 @@ def set_gpu_thread_mode_and_count(flags_obj): ...@@ -115,7 +115,7 @@ def set_gpu_thread_mode_and_count(flags_obj):
tf.compat.v1.logging.info('Logical CPU cores: %s', cpu_count) tf.compat.v1.logging.info('Logical CPU cores: %s', cpu_count)
# Allocate private thread pool for each GPU to schedule and launch kernels # Allocate private thread pool for each GPU to schedule and launch kernels
per_gpu_thread_count = flags_obj.per_gpu_thread_count or 2 per_gpu_thread_count = flags_obj.per_gpu_thread_count or 1
os.environ['TF_GPU_THREAD_MODE'] = flags_obj.tf_gpu_thread_mode os.environ['TF_GPU_THREAD_MODE'] = flags_obj.tf_gpu_thread_mode
os.environ['TF_GPU_THREAD_COUNT'] = str(per_gpu_thread_count) os.environ['TF_GPU_THREAD_COUNT'] = str(per_gpu_thread_count)
tf.compat.v1.logging.info('TF_GPU_THREAD_COUNT: %s', tf.compat.v1.logging.info('TF_GPU_THREAD_COUNT: %s',
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment