多GPU环境下在同一代码中分别运行Mask、Faster与YOLO v4的任务分配问题求助
解决多GPU任务拆分的问题
看起来你遇到的核心问题是模型初始化没有绑定到指定GPU,导致后续的tf.device上下文无法正确将任务分配到目标GPU。下面是具体的排查和修复方案:
1. 核心问题分析
你的代码里,segmentation_model(Mask分割任务)和Faster R-CNN的会话sess是在全局范围内初始化的,而非对应的tf.device上下文中。这会导致TensorFlow自动将模型参数加载到默认GPU(通常是GPU 0),后续即使切换tf.device,模型运算仍会在参数所在的GPU上执行,无法实现任务拆分。
另外,你当前用单while循环串行执行两个任务,即使设备分配正确,也无法利用多GPU并行处理,建议用线程实现并行运行。
2. 具体修复步骤
步骤1:在指定GPU上下文里初始化模型
把模型的创建、加载代码放到对应的tf.device块中,确保模型参数绑定到目标GPU:
import os os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"]="0,1" import pyzed.sl as sl import cv2 from pixellib.instance import instance_segmentation import time import threading import numpy as np import tensorflow as tf tf.debugging.set_log_device_placement(True) # 配置GPU内存增长 gpus_physical = tf.config.experimental.list_physical_devices('GPU') for gpu in gpus_physical: tf.config.experimental.set_memory_growth(gpu, True) gpus_logical = tf.config.list_logical_devices('GPU') from object_detection.utils import label_map_util from object_detection.utils import visualization_utils as vis_util import serial # --- 全局变量定义 --- frame = None # 共享摄像头帧数据 lock = threading.Lock() # 线程锁,避免帧数据竞争 # --- 任务1:Mask分割(绑定到GPU 1)--- def mask_task(): with tf.device(gpus_logical[1].name): # 在此初始化分割模型 segmentation_model = instance_segmentation() segmentation_model.load_model("你的Mask模型路径") # 替换为实际模型路径 while True: with lock: if frame is None: time.sleep(0.01) continue current_frame = frame.copy() res = segmentation_model.segmentFrame(current_frame, show_bboxes=False) output_frame = res[1] print("flag1") cv2.imshow("mask", output_frame) cv2.waitKey(1) # --- 任务2:Faster R-CNN检测(绑定到GPU 0)--- def faster_rcnn_task(): with tf.device(gpus_logical[0].name): # 在此初始化Faster R-CNN模型和会话 PATH_TO_CKPT = '你的Faster模型.pb路径' PATH_TO_LABELS = '你的label_map.pbtxt路径' NUM_CLASSES = 90 # 根据数据集调整 # 加载模型图 detection_graph = tf.Graph() with detection_graph.as_default(): od_graph_def = tf.GraphDef() with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name='') # 加载标签映射 label_map = label_map_util.load_labelmap(PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index = label_map_util.create_category_index(categories) # 创建会话 sess = tf.Session(graph=detection_graph) # 获取模型张量 image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') num_detections = detection_graph.get_tensor_by_name('num_detections:0') while True: with lock: if frame is None: time.sleep(0.01) continue current_frame = frame.copy() timer = time.time() image = current_frame image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image_expanded = np.expand_dims(image_rgb, axis=0) (boxes, scores, classes, num) = sess.run( [detection_boxes, detection_scores, detection_classes, num_detections], feed_dict={image_tensor: image_expanded}) # 后处理逻辑修正 conf_threshold = 0.5 nms_threshold = 0.45 image_shape = image.shape[:2] # 筛选高置信度结果 valid_indices = np.where(scores[0] > conf_threshold)[0] filtered_scores = scores[0][valid_indices] filtered_boxes = boxes[0][valid_indices] filtered_classes = classes[0][valid_indices] final_boxes = [] final_scores = [] final_classes = [] if len(filtered_scores) > 0: # 非极大值抑制(假设apply_non_max_suppression是你自定义的函数) indices, count = apply_non_max_suppression(filtered_boxes, filtered_scores, nms_threshold, 200) selected_indices = indices[:count] final_boxes = filtered_boxes[selected_indices] final_scores = filtered_scores[selected_indices] final_classes = filtered_classes[selected_indices] final_boxes = [denormalize_box(box, image_shape) for box in final_boxes] # 可视化检测结果 vis_util.visualize_boxes_and_labels_on_image_array( image, np.squeeze(boxes), np.squeeze(classes).astype(np.int32), np.squeeze(scores), category_index, use_normalized_coordinates=True, line_thickness=8, min_score_thresh=0.9) # 计算目标中心点 class_and_box = [] for i in range(len(final_boxes)): if final_scores[i]>0.90: case = {final_classes[i] : final_boxes[i]} class_and_box.append(case) def kutu_orta_nokta(liste): xmin, ymin, xmax, ymax = liste x, y = ((xmin+xmax)/2, (ymin+ymax)/2) center = int(x), int(y) cv2.circle(image, center, radius=10, color=(255, 0, 255), thickness=-1) return center sag_liste= [] orta_liste = [] sol_liste= [] cv2.imshow('Object detector', image) cv2.waitKey(1) print('[Info] Time Taken: {} | FPS: {}'.format(time.time() - timer, 1 / (time.time() - timer)), end='\r') # --- 主函数:启动摄像头和并行线程 --- if __name__ == "__main__": # 初始化ZED摄像头(根据实际情况调整参数) zed = sl.Camera() init_params = sl.InitParameters() init_params.camera_resolution = sl.RESOLUTION.HD720 init_params.camera_fps = 30 err = zed.open(init_params) if err != sl.ERROR_CODE.SUCCESS: print(f"摄像头打开失败:{err}") exit(1) # 启动两个任务线程 mask_thread = threading.Thread(target=mask_task) faster_thread = threading.Thread(target=faster_rcnn_task) mask_thread.daemon = True faster_thread.daemon = True mask_thread.start() faster_thread.start() # 主循环:读取摄像头帧 runtime_params = sl.RuntimeParameters() while True: if zed.grab(runtime_params) == sl.ERROR_CODE.SUCCESS: zed.retrieve_image(frame, sl.VIEW.LEFT) frame = frame.get_data() else: time.sleep(0.01) # 按q退出 if cv2.waitKey(1) & 0xFF == ord('q'): break # 清理资源 zed.close() cv2.destroyAllWindows()
步骤2:关键修复点说明
- 模型绑定GPU:将模型初始化、会话创建都放到对应
tf.device上下文,确保参数加载到目标GPU。 - 线程并行:用两个独立线程运行两个任务,让多GPU同时工作,提升效率。
- 线程安全:通过
threading.Lock保证摄像头帧数据的安全共享,避免数据竞争。 - 后处理修正:修复原代码中索引错误的问题,用numpy数组处理筛选结果更稳妥。
步骤3:验证GPU分配
运行代码后,用nvidia-smi查看GPU占用:
- GPU 0应主要被Faster R-CNN任务占用
- GPU 1应主要被Mask分割任务占用
如果仍有问题,可检查:
- 确认Pixellib的instance_segmentation支持TensorFlow设备上下文(默认支持)
- 查看TensorFlow的设备放置日志(已开启
tf.debugging.set_log_device_placement(True)),确认操作是否分配到指定GPU
3. YOLOv4任务的扩展
若要加入YOLOv4任务,只需新增一个线程,将YOLOv4的模型初始化放到对应的tf.device上下文(比如GPU 2,若有第三个GPU),逻辑和上述两个任务一致即可。
内容的提问来源于stack exchange,提问作者Ugur
相关产品推荐
相关产品推荐

