| 524 | |
| 525 | |
| 526 | def set_env_variables(gpu_id): |
| 527 | # Disable tensorflow logs, except in the case of an error. |
| 528 | if "TF_CPP_MIN_LOG_LEVEL" not in os.environ: |
| 529 | os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" |
| 530 | |
| 531 | # Set sharing strategy to file_system to avoid file descriptor leaks |
| 532 | torch.multiprocessing.set_sharing_strategy("file_system") |
| 533 | |
| 534 | # Only use one GPU, if we have one. |
| 535 | # For Tensorflow |
| 536 | # TODO: Using USE with `--parallel` raises similar issue as https://github.com/tensorflow/tensorflow/issues/38518# |
| 537 | os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) |
| 538 | # For PyTorch |
| 539 | torch.cuda.set_device(gpu_id) |
| 540 | |
| 541 | # Fix TensorFlow GPU memory growth |
| 542 | try: |
| 543 | import tensorflow as tf |
| 544 | |
| 545 | gpus = tf.config.experimental.list_physical_devices("GPU") |
| 546 | if gpus: |
| 547 | try: |
| 548 | # Currently, memory growth needs to be the same across GPUs |
| 549 | gpu = gpus[gpu_id] |
| 550 | tf.config.experimental.set_visible_devices(gpu, "GPU") |
| 551 | tf.config.experimental.set_memory_growth(gpu, True) |
| 552 | except RuntimeError as e: |
| 553 | print(e) |
| 554 | except ModuleNotFoundError: |
| 555 | pass |
| 556 | |
| 557 | |
| 558 | def attack_from_queue( |