diff --git a/app.py b/app.py index c317a09..c2915e5 100644 --- a/app.py +++ b/app.py @@ -22,7 +22,7 @@ # load in weights and classes physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: - tf.config.experimental.set_memory_growth(physical_devices[0], True) + tf.config.set_visible_devices(physical_devices[0:1], 'GPU') if tiny: yolo = YoloV3Tiny(classes=num_classes) diff --git a/detect.py b/detect.py index d0ade74..b3f3f2b 100644 --- a/detect.py +++ b/detect.py @@ -23,7 +23,7 @@ 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) + tf.config.set_visible_devices(physical_devices[0:1], 'GPU') if FLAGS.tiny: yolo = YoloV3Tiny(classes=FLAGS.num_classes) diff --git a/detect_video.py b/detect_video.py index f54de1f..553b383 100644 --- a/detect_video.py +++ b/detect_video.py @@ -25,7 +25,7 @@ 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) + tf.config.set_visible_devices(physical_devices[0:1], 'GPU') if FLAGS.tiny: yolo = YoloV3Tiny(classes=FLAGS.num_classes) diff --git a/yolov3_tf2/models.py b/yolov3_tf2/models.py index aa6d815..13ef549 100644 --- a/yolov3_tf2/models.py +++ b/yolov3_tf2/models.py @@ -208,7 +208,7 @@ def YoloV3(size=None, channels=3, anchors=yolo_anchors, masks=yolo_anchor_masks, classes=80, training=False): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: - tf.config.experimental.set_memory_growth(physical_devices[0], True) + tf.config.set_visible_devices(physical_devices[0:1], 'GPU') x = inputs = Input([size, size, channels], name='input') x_36, x_61, x = Darknet(name='yolo_darknet')(x) @@ -242,7 +242,7 @@ def YoloV3Tiny(size=None, channels=3, anchors=yolo_tiny_anchors, masks=yolo_tiny_anchor_masks, classes=80, training=False): physical_devices = tf.config.experimental.list_physical_devices('GPU') if len(physical_devices) > 0: - tf.config.experimental.set_memory_growth(physical_devices[0], True) + tf.config.set_visible_devices(physical_devices[0:1], 'GPU') x = inputs = Input([size, size, channels], name='input') x_8, x = DarknetTiny(name='yolo_darknet')(x)