多 GPU 推理张量流

Question

我想使用 tensorflow/Keras 执行多 GPU 推理

这是我的预测

 model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

 # Load weights trained on MS-COCO
 model.load_weights(COCO_MODEL_PATH, by_name=True)

 # COCO Class names
 # Index of the class in the list is its ID. For example, to get ID of
 # the teddy bear class, use: class_names.index('teddy bear')
 class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
                'bus', 'train', 'truck', 'boat', 'traffic light',
                'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
                'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
                'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
                'kite', 'baseball bat', 'baseball glove', 'skateboard',
                'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
                'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
                'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
                'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
                'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
                'teddy bear', 'hair drier', 'toothbrush']


 # Load a random image from the images folder
 file_names = next(os.walk(IMAGE_DIR))[2]
 image = skimage.io.imread(os.path.join(IMAGE_DIR, random.choice(file_names)))

 # Run detection
 results = model.detect([image], verbose=1)

 # Visualize results
 r = results[0]

有没有办法在多个 GPU 上运行这个模型？

提前致谢。

Answer 1

根据系统中的 GPU 数量增加 GPU_COUNT，并在使用modellib.MaskRCNN 创建模型时传递新的config。

class InferenceConfig(coco.CocoConfig):
    GPU_COUNT = 1 # increase the GPU count based on number of GPUs
    IMAGES_PER_GPU = 1

config = InferenceConfig()
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

https://github.com/matterport/Mask_RCNN/blob/master/samples/demo.ipynb