From 30ae6e08f4dba64006425e5587517712c1f97e8f Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 30 Aug 2021 03:31:32 -0400 Subject: [PATCH] set `allow_growth` in `default_tf_session_config` Currently, TF only initializes GPUs once, and never releases memory. As a result, we must set `allow_growth` in the first Session (or every session?), otherwise it doesn't work. Please confirm the `allow_growth` config in `trainer.py` is expected, otherwise we should remove it instead. --- deepmd/env.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/env.py b/deepmd/env.py index aaa148b357..4868eb8c3b 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -119,6 +119,7 @@ def get_tf_session_config() -> Any: set_tf_default_nthreads() intra, inter = get_tf_default_nthreads() config = tf.ConfigProto( + gpu_options=tf.GPUOptions(allow_growth=True), intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter ) return config