Skip to content

GitLab

"docs/git@developer.sourcefind.cn:OpenDAS/ollama.git" did not exist on "d77c1c5f9d619de9544d726da265a18edd6b0692"
Unverified Commit ece8deb6 authored by Haoyu Zhang's avatar Haoyu Zhang Committed by GitHub
Browse files

Set default GPU private threads to 2 per device (#6524) 

This should give better performance in tweaked benchmarks.
parent 17e923da
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
official/resnet/keras/keras_common.py
View file @ ece8deb6
...@@ -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 1 per_gpu_thread_count = flags_obj.per_gpu_thread_count or 2
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