DeepRegNet /
DeepReg
| @@ 220-232 (lines=13) @@ | ||
| 217 | # env vars |
|
| 218 | os.environ["CUDA_VISIBLE_DEVICES"] = gpu |
|
| 219 | os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "false" if gpu_allow_growth else "true" |
|
| 220 | if num_workers <= 0: # pragma: no cover |
|
| 221 | logger.info( |
|
| 222 | "Limiting CPU usage by setting environment variables " |
|
| 223 | "OMP_NUM_THREADS, TF_NUM_INTRAOP_THREADS, TF_NUM_INTEROP_THREADS to %d. " |
|
| 224 | "This may slow down the prediction. " |
|
| 225 | "Please use --num_workers flag to modify the behavior. " |
|
| 226 | "Setting to 0 or negative values will remove the limitation.", |
|
| 227 | num_workers, |
|
| 228 | ) |
|
| 229 | # limit CPU usage |
|
| 230 | # https://github.com/tensorflow/tensorflow/issues/29968#issuecomment-789604232 |
|
| 231 | os.environ["OMP_NUM_THREADS"] = str(num_workers) |
|
| 232 | os.environ["TF_NUM_INTRAOP_THREADS"] = str(num_workers) |
|
| 233 | os.environ["TF_NUM_INTEROP_THREADS"] = str(num_workers) |
|
| 234 | ||
| 235 | # load config |
|
| @@ 89-101 (lines=13) @@ | ||
| 86 | # set env variables |
|
| 87 | os.environ["CUDA_VISIBLE_DEVICES"] = gpu |
|
| 88 | os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true" if gpu_allow_growth else "false" |
|
| 89 | if num_workers <= 0: # pragma: no cover |
|
| 90 | logger.info( |
|
| 91 | "Limiting CPU usage by setting environment variables " |
|
| 92 | "OMP_NUM_THREADS, TF_NUM_INTRAOP_THREADS, TF_NUM_INTEROP_THREADS to %d. " |
|
| 93 | "This may slow down the training. " |
|
| 94 | "Please use --num_workers flag to modify the behavior. " |
|
| 95 | "Setting to 0 or negative values will remove the limitation.", |
|
| 96 | num_workers, |
|
| 97 | ) |
|
| 98 | # limit CPU usage |
|
| 99 | # https://github.com/tensorflow/tensorflow/issues/29968#issuecomment-789604232 |
|
| 100 | os.environ["OMP_NUM_THREADS"] = str(num_workers) |
|
| 101 | os.environ["TF_NUM_INTRAOP_THREADS"] = str(num_workers) |
|
| 102 | os.environ["TF_NUM_INTEROP_THREADS"] = str(num_workers) |
|
| 103 | ||
| 104 | # load config |
|
