Merge pull request zzh8829#85 from friyin/master

Fix for Nvidia RTX 2xxx series
zcdliuwei · Nov 14, 2019 · 4cd9a9e · 4cd9a9e
2 parents 4f8a3c7 + 05dc413
commit 4cd9a9e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 0 deletions.
diff --git a/detect.py b/detect.py
@@ -21,6 +21,10 @@
 
 
 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)
+
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)
     else:

diff --git a/detect_video.py b/detect_video.py
@@ -23,6 +23,10 @@
 
 
 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)
+
     if FLAGS.tiny:
         yolo = YoloV3Tiny(classes=FLAGS.num_classes)
     else:

diff --git a/train.py b/train.py
@@ -42,6 +42,10 @@
 
 
 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)
+
     if FLAGS.tiny:
         model = YoloV3Tiny(FLAGS.size, training=True,
                            classes=FLAGS.num_classes)