Import Geant4 11.4.0.beta source tree

examples/extended/parameterisations/Par04/training_vae/utils/gpu_limiter.py

@@ -0,0 +1,36 @@
+import os
+from dataclasses import dataclass
+
+import tensorflow as tf
+
+
+@dataclass
+class GPULimiter:
+    """
+    Class responsible to set the limits of possible GPU usage by TensorFlow. Currently, the limiter creates one
+    instance of logical device per physical device. This can be changed in a future.
+
+    Attributes:
+        _gpu_ids: A string representing visible devices for the process. Identifiers of physical GPUs should
+            be separated by commas (no spaces).
+        _max_gpu_memory_allocation: An integer specifying limit of allocated memory per logical device.
+
+    """
+    _gpu_ids: str
+    _max_gpu_memory_allocation: int
+
+    def __call__(self):
+        os.environ["CUDA_VISIBLE_DEVICES"] = f"{self._gpu_ids}"
+        gpus = tf.config.list_physical_devices('GPU')
+        if gpus:
+            # Restrict TensorFlow to only allocate max_gpu_memory_allocation*1024 MB of memory on one of the GPUs
+            try:
+                for gpu in gpus:
+                    tf.config.set_logical_device_configuration(
+                        gpu,
+                        [tf.config.LogicalDeviceConfiguration(memory_limit=1024 * self._max_gpu_memory_allocation)])
+                logical_gpus = tf.config.list_logical_devices('GPU')
+                print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+            except RuntimeError as e:
+                # Virtual devices must be set before GPUs have been initialized
+                print(e)