Yolov5 Torchscript + Batch + 多GPU推理

Posted 2021-10-03 洪流之源

1. 导出torchscript模型

在yolov5工程目录下执行如下命令导出torchscript模型

python models/export.py --weights weights/yolov5m.pt

2. Yolov5推理类

class Yolov5Inference:

    def __init__(self, model_file, device, 
                 input_size_h, input_size_w, 
                 class_label, anchors=None,
                 obj_conf_thresh=0.3, nms_iou_thresh=0.5
                 ):

        self.model_file = model_file
        self.device = device
        self.input_size_h = input_size_h
        self.input_size_w = input_size_w
        self.obj_conf_thresh = obj_conf_thresh
        self.nms_iou_thresh = nms_iou_thresh
        self.class_label = class_label

        if anchors is None:
            self.anchors = [[10,13], [16,30], [33,23],
                            [30,61], [62,45], [59,119],
                            [116,90],[156,198], [373,326]]
        else:
            self.anchors = anchors

        s