1. Открытый исходный код github

Открытый исходный код находится здесь , код модели darknet для ONNX основан на Python, а код вывода TensorRT основан на C ++.

2. Преобразование модели даркнета в модель ONNX.

Модель даркнет может быть преобразована в модель ONNX через export_onnx.py файл, и он в настоящее время поддерживает YOLOv3, YOLOv3-SPP, YOLOv4 и другие модели.

    parser = argparse.ArgumentParser(description='Transform YOLO weights to ONNX.')
    parser.add_argument('--cfg_file', type=str, default='yolov4.cfg', help='yolo cfg file')
    parser.add_argument('--weights_file', type=str, default='yolov4.weights', help='yolo weights file')
    parser.add_argument('--output_file', type=str, default='yolov4.onnx', help='yolo onnx file')
    parser.add_argument('--neck', type=str, default='PAN', help='use which kind neck')
    args = parser.parse_args()

cfg_fileЭто путь к cfg-файлу darknet, путь weights_fileк файлу весов output_filedarknet, имя выходного файла ONNX, neckструктура модели YOLO, входных FPN моделей, связанных с YOLOv3, и входных PAN, связанных с моделями YOLOv4. После ввода следующей команды вы можете преобразовать модель даркнета в модель ONNX.

python3 export_onnx.py

Команда для экспорта модели yolov4-tiny:

python3 export_onnx.py --cfg_file yolov4-tiny.cfg --weights_file yolov4-tiny.weights --output_file yolov4-tiny.onnx --strides 32 16 --neck FPN

Транспонирование добавляется к преобразованной модели, и выходные результаты всех слоев YOLO объединяются. Окончательная структура выглядит следующим образом:
Вставьте описание изображения сюда

3. Преобразование модели ONNX в модель TensorRT.

3.1 Обзор

Модель TensorRT - это механизм рассуждений TensorRT, а код реализован на C ++. Соответствующая конфигурация записана в файле config.yaml , если engine_fileпуть существует , он будет прочитан engine_file, иначе он будет onnx_fileсгенерирован engine_file.

void YOLOv5::LoadEngine() {
    // create and load engine
    std::fstream existEngine;
    existEngine.open(engine_file, std::ios::in);
    if (existEngine) {
        readTrtFile(engine_file, engine);
        assert(engine != nullptr);
    } else {
        onnxToTRTModel(onnx_file, engine_file, engine, BATCH_SIZE);
        assert(engine != nullptr);
    }
}

В файле config.yaml можно установить размер пакета, размер изображения и привязку модели.

Yolov4:
    onnx_file:     "../yolov4.onnx"
    engine_file:   "../yolov4.trt"
    labels_file:   "../coco.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   608
    IMAGE_HEIGHT:  608
    obj_threshold: 0.4
    nms_threshold: 0.45
    stride:        [8, 16, 32]
    anchors:       [[12, 16], [19, 36], [40, 28], [36, 75], [76, 55], [72, 146], [142, 110], [192, 243], [459, 401]]

3.2 Компиляция

Скомпилируйте проект с помощью следующей команды, чтобы сгенерироватьYolov4_trt

mkdir build && cd build
cmake ..
make -j

3.3 Эксплуатация

Запустите проект с помощью следующей команды, чтобы получить результат вывода

./Yolov4_trt ../config.yaml ../samples

4. Результаты вывода

Результат вывода показан на рисунке ниже:
Вставьте описание изображения сюда
вы можете видеть, что выходной результат точно такой же, как выходной результат модели даркнета, включая уровень достоверности. В выходной информации среднее время обработки каждого изображения составляет около 16 мс. , а время вывода одного двигателя составляет около 7 мс.

Ubuntu18.04 настраивает среду даркнета для достижения обнаружения цели YOLOv4 (5) - модели Darknet YOLOv4 и YOLOv4-tiny превращают ONNX в развертывание TensorRT

Каталог статей