This is an automated email from the ASF dual-hosted git repository. domino pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/madlib.git
commit 229d2b1cac9b19cee298f96ebf3c28818a0cd038 Author: Domino Valdano <dvald...@vmware.com> AuthorDate: Wed Nov 11 23:22:25 2020 -0800 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. --- .../postgres/modules/deep_learning/madlib_keras_wrapper.py_in | 8 ++++++++ 1 file changed, 8 insertions(+) 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 @@ def reset_cuda_env(value): 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 @@ def get_keras_session(device_name, gpu_count, segments_per_host): 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):