| Geant4 Cross Reference |
1 import os 1 import os
2 from dataclasses import dataclass 2 from dataclasses import dataclass
3 3
4 import tensorflow as tf 4 import tensorflow as tf
5 5
6 6
7 @dataclass 7 @dataclass
8 class GPULimiter: 8 class GPULimiter:
9 """ 9 """
10 Class responsible to set the limits of pos 10 Class responsible to set the limits of possible GPU usage by TensorFlow. Currently, the limiter creates one
11 instance of logical device per physical de 11 instance of logical device per physical device. This can be changed in a future.
12 12
13 Attributes: 13 Attributes:
14 _gpu_ids: A string representing visibl 14 _gpu_ids: A string representing visible devices for the process. Identifiers of physical GPUs should
15 be separated by commas (no spaces) 15 be separated by commas (no spaces).
16 _max_gpu_memory_allocation: An integer 16 _max_gpu_memory_allocation: An integer specifying limit of allocated memory per logical device.
17 17
18 """ 18 """
19 _gpu_ids: str 19 _gpu_ids: str
20 _max_gpu_memory_allocation: int 20 _max_gpu_memory_allocation: int
21 21
22 def __call__(self): 22 def __call__(self):
23 os.environ["CUDA_VISIBLE_DEVICES"] = f 23 os.environ["CUDA_VISIBLE_DEVICES"] = f"{self._gpu_ids}"
24 gpus = tf.config.list_physical_devices 24 gpus = tf.config.list_physical_devices('GPU')
25 if gpus: 25 if gpus:
26 # Restrict TensorFlow to only allo 26 # Restrict TensorFlow to only allocate max_gpu_memory_allocation*1024 MB of memory on one of the GPUs
27 try: 27 try:
28 for gpu in gpus: 28 for gpu in gpus:
29 tf.config.set_logical_devi 29 tf.config.set_logical_device_configuration(
30 gpu, 30 gpu,
31 [tf.config.LogicalDevi 31 [tf.config.LogicalDeviceConfiguration(memory_limit=1024 * self._max_gpu_memory_allocation)])
32 logical_gpus = tf.config.list_ 32 logical_gpus = tf.config.list_logical_devices('GPU')
33 print(len(gpus), "Physical GPU 33 print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
34 except RuntimeError as e: 34 except RuntimeError as e:
35 # Virtual devices must be set 35 # Virtual devices must be set before GPUs have been initialized
36 print(e) 36 print(e)