From fad33a4ec17b3796987b7d7ad5b7f1a9956fecd4 Mon Sep 17 00:00:00 2001 From: Trijeet Modak Date: Fri, 5 Mar 2021 11:52:11 +0530 Subject: [PATCH] added '--allow-growth' flag added '--allow-growth' flag to choose to limit GPU growth. This solves `failed to create cublas handle: CUBLAS_STATUS_NOT_INITIALIZED` issue if caused due to GPU memory limit. #3502 --- research/object_detection/model_main_tf2.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/research/object_detection/model_main_tf2.py b/research/object_detection/model_main_tf2.py index 0cf053039ec..9accb76d299 100644 --- a/research/object_detection/model_main_tf2.py +++ b/research/object_detection/model_main_tf2.py @@ -68,6 +68,9 @@ ('Whether or not to record summaries during' ' training.')) +flags.DEFINE_boolean('allow_growth', False, + ('Whether or not to limit gpu memory growth.')) + FLAGS = flags.FLAGS @@ -75,6 +78,17 @@ def main(unused_argv): flags.mark_flag_as_required('model_dir') flags.mark_flag_as_required('pipeline_config_path') tf.config.set_soft_device_placement(True) + + if FLAGS.allow_growth: + gpus = tf.config.list_physical_devices('GPU') + if gpus: + try: + # Currently, memory growth needs to be the same across GPUs + for gpu in gpus: + tf.config.experimental.set_memory_growth(gpu, True) + except RuntimeError as e: + # Memory growth must be set before GPUs have been initialized + pass if FLAGS.checkpoint_dir: model_lib_v2.eval_continuously(