Gabriele Cosmo
37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
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)
|