Fix issue when setting memory growth

When calling tf.config.experimental.set_memory_growth, we want to set the value to True for all physical devices. Otherwise the following ValueError is produce: 'Memory growth cannot differ between GPU devices'
BUMETCSDA · Apr 5, 2020 · 67d666f · 67d666f
1 parent 65294d5
commit 67d666f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/detect.py b/detect.py
@@ -23,8 +23,8 @@
 
 def main(_argv):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
-    if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+    for physical_device in physical_devices:
+        tf.config.experimental.set_memory_growth(physical_device, True)
 
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)

diff --git a/detect_video.py b/detect_video.py
@@ -24,8 +24,8 @@
 
 def main(_argv):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
-    if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+    for physical_device in physical_devices:
+        tf.config.experimental.set_memory_growth(physical_device, True)
 
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)
@@ -63,7 +63,7 @@ def main(_argv):
             time.sleep(0.1)
             continue
 
-        img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 
+        img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         img_in = tf.expand_dims(img_in, 0)
         img_in = transform_images(img_in, FLAGS.size)
 

diff --git a/train.py b/train.py
@@ -46,8 +46,8 @@
 
 def main(_argv):
     physical_devices = tf.config.experimental.list_physical_devices('GPU')
-    if len(physical_devices) > 0:
-        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+    for physical_device in physical_devices:
+        tf.config.experimental.set_memory_growth(physical_device, True)
 
     if FLAGS.tiny:
         model = YoloV3Tiny(FLAGS.size, training=True,