You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

多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

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.04.27 19:57:49