yolov5PyTorchモデルをTensorRTに変換します

1.githubオープンソースコード

yolov5 TensorRT推論のオープンソースコードの場所はhttps://github.com/linghu8812/tensorrt_inference/tree/master/yolov5で、PyTorchからonnxへのコードは、元の作成者のフォークからのコードにあります：https：/ /github.com/linghu8812/ yolov5、モデル変換に少し変更を加えました。

2.PyTorchモデルをONNXモデルに変換します

最初にコマンドgit clone https://github.com/linghu8812/yolov5.gitcloneyolov5 codeを使用し、次にexport_onnx.pyファイルをyolov5/modelsフォルダーにコピーし、次のコマンドを使用してONNXファイルを生成します。yolov5s、yolov5m、yolov5l、yolov5xの場合、これらのモデルはすべてサポートできます。--weightsモデルファイルのパスを指定--imgし、入力画像--batchサイズのバッチサイズを設定できます。

export PYTHONPATH="$PWD" && python3 export_onnx.py --weights ./weights/yolov5s.pt --img 640 --batch 1

YOLOv4の図に示すようにモデル、連結の出力も、作ら。onnxsimモジュール
yolov5
がコードに追加され、後続のトランスポーズノードとリシェイプノードの構造が簡略化されます。簡略化されたコードは次のとおりです。

        onnx_model = onnx.load(f)  # load onnx model
        model_simp, check = simplify(onnx_model)
        assert check, "Simplified ONNX model could not be validated"
        onnx.save(model_simp, f)

モデル構造は次のように簡略化されます。
ここに画像の説明を挿入します

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ファイルは、バッチサイズ、画像のサイズ、モデルのアンカーを設定できます。

yolov5:
    onnx_file:     "../yolov5x.onnx"
    engine_file:   "../yolov5x.trt"
    labels_file:   "../coco.names"
    BATCH_SIZE:    1
    INPUT_CHANNEL: 3
    IMAGE_WIDTH:   640
    IMAGE_HEIGHT:  640
    obj_threshold: 0.4
    nms_threshold: 0.45
    stride:        [8, 16, 32]
    anchors:       [[10,13], [16,30], [33,23], [30,61], [62,45], [59,119], [116,90], [156,198], [373,326]]

3.2コンパイル

次のコマンドを使用してプロジェクトをコンパイルし、生成しますyolov5_trt

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

3.3操作

次のコマンドでプロジェクトを実行して、推論結果を取得します

./yolov5_trt ../config.yaml ../samples

4.推論結果

推論結果を次の図に示します。
ここに画像の説明を挿入します

上の写真はyolov5xのテスト結果です。出力情報では、yolov5xの各写真の平均処理時間は約22.9ms、単一エンジンの推論時間は約12.1msです。