gpu_utils.py
1.25 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import os
def tf_setup(GPU_MEM_LIMIT):
os.environ["TF_USE_LEGACY_KERAS"] = "1"
if GPU_MEM_LIMIT == 0:
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import tensorflow as tf
from transformers import __version__ as transformers_version
try:
import tf_keras as keras
except:
import keras
gpus = tf.config.list_physical_devices('GPU')
# https://www.tensorflow.org/guide/gpu
if gpus:
try:
# Currently, memory growth needs to be the same across GPUs
#for gpu in gpus:
# tf.config.experimental.set_memory_growth(gpu, True)
tf.config.set_logical_device_configuration(
gpus[0],
[tf.config.LogicalDeviceConfiguration(memory_limit=GPU_MEM_LIMIT)]
)
logical_gpus = tf.config.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
except RuntimeError as e:
# Memory growth must be set before GPUs have been initialized
print(e)
print(tf.test.gpu_device_name())
print('TENSORFLOW: ', tf.__version__)
print('KERAS: ', keras.__version__)
print('TRANSFORMERS:', transformers_version)