taco/microWakeWord-Trainer-Nvidia-Docker
1
0
Fork 0
mirror of https://github.com/TaterTotterson/microWakeWord-Trainer-Nvidia-Docker.git synced 2026-06-12 20:10:19 -06:00
Code Issues Packages Projects Releases Wiki Activity

Improve CUDA/XLA setup and CPU fallback

Browse Source
This commit is contained in:
Tater Totterson
2026-01-20 20:47:58 -06:00
parent e72773f0ef
commit 8669ccdfd4
4 changed files with 38 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
train_wake_word
View File

@@ -70,7 +70,10 @@ START_TS=$EPOCHSECONDS
export TF_CPP_MIN_LOG_LEVEL=9
export TF_FORCE_GPU_ALLOW_GROWTH=true
export TF_GPU_ALLOCATOR=cuda_malloc_async
export TF_XLA_FLAGS="--tf_xla_auto_jit=0"
DEFAULT_XLA_FLAGS="--tf_xla_auto_jit=0 --xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found --xla_gpu_cuda_data_dir=${DATA_DIR}/cuda"
DEFAULT_XLA_RUNTIME_FLAGS="--xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found --xla_gpu_cuda_data_dir=${DATA_DIR}/cuda"
export TF_XLA_FLAGS="${TF_XLA_FLAGS:+${TF_XLA_FLAGS} }${DEFAULT_XLA_FLAGS}"
export XLA_FLAGS="${XLA_FLAGS:+${XLA_FLAGS} }${DEFAULT_XLA_RUNTIME_FLAGS}"
export NVIDIA_TF32_OVERRIDE=1
export TF_CUDNN_WORKSPACE_LIMIT_IN_MB=512
export GLOG_minloglevel=2
@@ -127,4 +130,4 @@ print_elapsed_time --no-separators "${POST_GEN_TS}" "${POST_AUGMENT_TS}" "Augmen
print_elapsed_time --no-separators "${POST_AUGMENT_TS}" "${END_TS}" "${TRAINING_STEPS} training steps"
python -c $'msg="="*54 ; print(f"{msg:>80s}")'
print_elapsed_time --no-separators "${START_TS}" "${END_TS}" "Total"
python -c $'print(f"{\'=\' * 80}")'
python -c $'print(f"{\'=\' * 80}")'