Question

我在imageai中使用Video object detection python模块。我创建了一个flask app并使用threading进行并行操作。

我正在尝试将Video object detection工作（http请求）统一分发到可用的GPUs（例如，我正在随机分发）下面提供了worker函数。但仍然仅利用GPU 0和我的concurrent http requests停留在队列中（依次执行）

PS_OPS = ['Variable', 'VariableV2', 'AutoReloadVariable']

def assign_to_device(device, ps_device='/cpu:0'):
    def _assign(op):
        node_def = op if isinstance(op, tf.NodeDef) else op.node_def
        if node_def.op in PS_OPS:
            return "/" + ps_device
        else:
            return device
    return _assign


def generic_worker(h5, model_type,type):
    def worker(iq, oq):
        if(type=="prediction"):
            prediction = ImagePrediction()
        else:
            prediction = VideoObjectDetection()
        getattr(prediction, model_type)()
        prediction.setModelPath(os.path.join(app.root_path,"models", h5))
        prediction.loadModel()

        while True:
            import tensorflow as tf
            rid = iq.get()
            mylist =[]
            gpu_list = ['0','1','2','3']
            i=random.choice(gpu_list)
            with tf.device(assign_to_device('/gpu:{}'.format(i), ps_device='/cpu:0')): ## Tried to execute worker in random GPU
                print("Now in GPU ",i)
                video_files = os.listdir(os.path.join(app.root_path ,"detection_videos"))
                    video_name=random.choice(video_files)
                    video_path = prediction.detectObjectsFromVideo(input_file_path=os.path.join(app.root_path,"detection_videos", video_name),
                        output_file_path=os.path.join(app.root_path,"detected_videos", video_name)
                        , frames_per_second=20, log_progress=False)
                    return_dict = {}                  
                    return_dict["GPU_"+str(i)+"_Input"] = os.path.join(app.root_path,"detection_videos", video_name)
                    return_dict["GPU_"+str(i)+"_Output"] = video_path

Imageai多GPU支持

0 个答案: