DL: Enable JIT XLA auto-clustering, if available.
This is a simple optimization for CPU and/or GPU we can turn on,
to improve performance.
If it's not available, that means tensorflow version is too old.
Will give a warning, suggesting to upgrade to 1.14.0 and continue
without it.
diff --git a/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in b/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
index d7b2d41..7b1b3de 100644
--- a/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
+++ b/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
@@ -65,6 +65,13 @@
if CUDA_VISIBLE_DEVICES_KEY in os.environ:
del os.environ[CUDA_VISIBLE_DEVICES_KEY]
+def enable_xla():
+ os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit'
+ try:
+ tf.config.optimizer.set_jit(True)
+ except:
+ plpy.warning("This version of tensorflow does not fully support XLA auto-cluster JIT optimization. HINT: upgrading to tensorflow 1.14.0 may improve performance.")
+
def get_device_name_and_set_cuda_env(gpu_count, seg):
if gpu_count > 0:
device_name = '/gpu:0'
@@ -90,6 +97,7 @@
config.gpu_options.allow_growth = False
config.gpu_options.per_process_gpu_memory_fraction = memory_fraction
session = tf.Session(config=config)
+ enable_xla()
return session
def clear_keras_session(sess = None):
