基于Ascend910+PyTorch1.11.0+CANN6.3.RC2的YoloV5训练推理一体化解决方案

昇腾Pytorch镜像：https://ascendhub.huawei.com/#/detail/ascend-pytorch

代码仓：git clone https://gitee.com/ascend/modelzoo-GPL.git

coco测试验证集：wget https://bj-aicc.obs.cn-north-309.mtgascendic.cn/dataset/coco2017/coco.zip

coco训练集(放images下）：wget https://bj-aicc.obs.cn-north-309.mtgascendic.cn/dataset/coco2017/train2017.zip

部分代码

# import StreamManagerApi.py
from StreamManagerApi import *

if __name__ == '__main__':
    # init stream manager
    streamManagerApi = StreamManagerApi()
    ret = streamManagerApi.InitManager()
    if ret != 0:
        print("Failed to init Stream manager, ret=%s" % str(ret))
        exit()

    # create streams by pipeline config file
    with open("data/pipeline/Sample.pipeline", 'rb') as f:
        pipelineStr = f.read()
    ret = streamManagerApi.CreateMultipleStreams(pipelineStr)
    if ret != 0:
        print("Failed to create Stream, ret=%s" % str(ret))
        exit()

    # Construct the input of the stream
    dataInput = MxDataInput()
    with open("data/test.jpg", 'rb') as f:
        dataInput.data = f.read()

    # The following is how to set the dataInput.roiBoxs
    """
    roiVector = RoiBoxVector()
    roi = RoiBox()
    roi.x0 = 100
    roi.y0 = 100
    roi.x1 = 200
    roi.y1 = 200
    roiVector.push_back(roi)
    dataInput.roiBoxs = roiVector
    """

    # Inputs data to a specified stream based on streamName.
    streamName = b'classification'
    inPluginId = 0
    uniqueId = streamManagerApi.SendDataWithUniqueId(streamName, inPluginId, dataInput)
    if uniqueId < 0:
        print("Failed to send data to stream.")
        exit()

    # Obtain the inference result by specifying streamName and uniqueId.
    inferResult = streamManagerApi.GetResultWithUniqueId(streamName, uniqueId, 3000)
    if inferResult.errorCode != 0:
        print("GetResultWithUniqueId error. errorCode=%d, errorMsg=%s" % (
            inferResult.errorCode, inferResult.data.decode()))
        exit()

    # print the infer result
    print(inferResult.data.decode())

    # destroy streams
    streamManagerApi.DestroyAllStreams()

本来想一次性写完的，奈何装CANN的驱动装了一个礼拜，各种内核版本不匹配，国产AI硬件任重道远...