From 26cbdeb683c4898e48a5c603bfc7bf7ff3219058 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 24 May 2022 23:21:03 -0400 Subject: [PATCH] add JIT option --- deepmd/env.py | 2 ++ doc/train/training-advanced.md | 1 + 2 files changed, 3 insertions(+) diff --git a/deepmd/env.py b/deepmd/env.py index db19fd86c6..8e2d964885 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -181,6 +181,8 @@ def get_tf_session_config() -> Any: """ set_tf_default_nthreads() intra, inter = get_tf_default_nthreads() + if int(os.environ.get("DP_JIT", 0)): + set_env_if_empty("TF_XLA_FLAGS", "--tf_xla_auto_jit=2") config = tf.ConfigProto( gpu_options=tf.GPUOptions(allow_growth=True), intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 74998f82a7..ca0269db67 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -155,6 +155,7 @@ One can set other environmental variables: | --------------------- | ---------------------- | ------------- | -------------------------- | | DP_INTERFACE_PREC | `high`, `low` | `high` | Control high (double) or low (float) precision of training. | | DP_AUTO_PARALLELIZATION | 0, 1 | 0 | Enable auto parallelization for CPU operators. | +| DP_JIT | 0, 1 | 0 | Enable JIT. Note that this option may either improve or decrease the performance. Requires TensorFlow supports JIT. | ## Adjust `sel` of a frozen model