tf-keras_header.py

import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '0' # -1 = no GPU

import tensorflow as tf
import tensorflow.keras as keras

# versions
print('tf version: {}'.format(tf.__version__))
print('keras version: {}'.format(keras.__version__))
print('build with CUDA:', tf.test.is_built_with_cuda())

# disable Eager execution
# non-eager (graph mode): faster execution; eager: easier debugging
tf.compat.v1.disable_eager_execution()
tf.compat.v1.experimental.output_all_intermediates(True)

# choose GPU and limit GPU memory usage
gpus = tf.config.list_physical_devices('GPU')
gpu_id = 0

if gpus:
    try:
        tf.config.set_visible_devices(gpus[gpu_id], 'GPU')
        tf.config.experimental.set_memory_growth(gpus[gpu_id], True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(len(gpus), 'Physical GPU(s),', len(logical_gpus), 'Logical GPU(s)')
    except RuntimeError as e:
        print(e)
else:
    print('No GPU found!')

# for TF <= 2.0.0
# gpus = tf.config.experimental.list_physical_devices('GPU')
# gpu_id = 0 # to find GPU IDs and names: lookup through task manager, or run `nvidia-smi -L`
# if gpus:
#     try:
#         tf.config.experimental.set_visible_devices(gpus[gpu_id], 'GPU')
#         tf.config.experimental.set_memory_growth(gpus[gpu_id], True)
#         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
#         print(len(gpus), 'Physical GPU(s),', len(logical_gpus), 'Logical GPU(s)')
#     except RuntimeError as e:
#         print(e)

# disable Eager execution i.e. use graph mode
# graph: faster execution; eager: easier debugging
tf.compat.v1.disable_eager_execution()
tf.compat.v1.experimental.output_all_intermediates(True)