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