JetsonNano2GB編譯Yolov5並用TensorRT加速教程

小肉包老師發表於2020-11-14
NaN編譯YOLO

上一期我們教大家如何給新的JetsonNano 2GB燒錄系統。這一期我們將教大家如何在JetsonNano上部署最新的Yolov5檢測模型，並且採用TensorRT加速，看看我們的模型能否在JetsonNano這樣的小裝置上跑到實時。

image-20201107222319190.png

首先我們來確認一下系統裡面opencv對應的版本：

image-20201107223801733.png

是最新的4.1，不錯。

今天這是我們要實現的效果：

[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片儲存下來直接上傳(img-Us5L9WaW-1605327260590)(https://i.loli.net/2020/11/11/OeuA2y3MkF4fSbP.png)]

image.png

是的，你沒有看錯，我們將在JetsonNano上跑一個經過TensorRT加速的YoloV5演算法，演算法的功能是實現口罩佩戴的檢測。從圖中你可以看到，對於口罩佩戴的檢測還是很準確的，而且基本上可以達到實時！這裡要說明的是，這個模型的輸入尺寸不大，並且在Yolov5的基礎上經過了裁剪，所以能夠做到這麼快，但是如果你的模型直接拿官方的部署，可能無法跑到這個速度，最後速度帶來的當然是精度的降低，不過在可以接受的範圍之內！

還等什麼，快來依照我們的教程來實現吧~

技術棧

本篇文章不涉及到Python或者任何訓練方面的東西，純部署，而且純C++，因此教程裡面可能會出現一些你不熟悉的庫或者術語，沒關心，你看到最後我們跑出來的結果就可以了，至於過程你可以通過加入我們的社群來獲取想要的材料。大部分程式碼都在github開源。核心技術包括:

  • C++;
  • TensorRT相關基礎知識;
  • 對ONNX瞭解;
  • 對Yolov5的基本網路結構瞭解。

準備工作

先做一些準備工作，什麼準備工作呢？我們需要先安裝一個C++下面比較常用的視覺化庫，thor:

git clone https://github.com/jinfagang/thor

大家可以進入到主頁依照readme進行編譯。為什麼要做這個準備幹工作？原因有兩個：

  • 它依賴一些庫，如果你能夠成功編譯它，說明你的環境沒有問題，如果不能，你的環境可能有問題了；
  • 它可以很方便的視覺化我們的檢測輸出結果。

[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片儲存下來直接上傳(img-MmRVptpR-1605327260592)(https://i.loli.net/2020/11/11/mHyaS6R2q4fJpDk.png)]

引用一張thor官方的圖：

畫檢測結果既好看又方便！

編譯onnx-tensorrt

我們給大家準備好了一個轉好的onnx檔案，如果你需要匯出自己的onnx，過程有點複雜，你可以通過加入社群交流如何匯出，我們沒有現成的程式碼提供。

git clone --recursive https://github.com/onnx/onnx-tensorrt

編譯onnx-tensorrt 在JetsonNano上可能會遇到一些問題。為什麼呢？

第一個錯誤：

CMake Error at CMakeLists.txt:21 (cmake_minimum_required):
  CMake 3.13 or higher is required.  You are running version 3.10.2


-- Configuring incomplete, errors occurred!

然後你就需要重新編譯一下CMake了：

git clone https://github.com/Kitware/CMake

說實話這個CMake也是異常的龐大了。clone下來100多M。然後開始編譯CMake了：

./bootstrap.sh
make
sudo make install
cmake -V

編譯過程還是挺慢的.

然後就可以接著編譯onnx-tensorrt了。

編譯完onnx-tensorrt之後，就可以用 onnx2trt 將模型轉出到tensorrt 的engine了。

在JetsonNano上將ONNX轉到engine檔案

想必大家已經安裝好了必備的工具，如果你在安裝工具的時候遇到問題，不要灰心喪氣，因為最終是可以work的。我們已經走到這一步了。

接下來就是拿一個onnx來做試驗了。這裡我推薦一個onnx檔案，大小隻有6M，我們將用它來實現前面的檢測效果。

至於如何訓練對應的模型，大家可以在神力平臺找到對應的訓練程式碼。

http://manaai.cn

我們成功的將onnx轉到了engine:

onnx2trt best_sim.onnx -o a.trt -w 10909920

[外鏈圖片轉存失敗,源站可能有防盜鏈機制,建議將圖片儲存下來直接上傳(img-FlBdC4Ex-1605327260593)(https://i.loli.net/2020/11/14/kCdPOfblJW41cEz.png)]

這裡需要注意的是這個-w, 它指的是我們的workspace多大，你不要預設，也不要設定的太小，大概按照我這樣的設定就可以了，預設的話就當機了，太小會報錯。

對TensorRT Engine進行推理

接下來就要對這個Engine檔案進行推理了。推理部分我們要實現load上面的tensorrt engine並實現口罩的後處理檢測。對應的程式碼也可以在神力平臺找到：

http://manaai.cn

最後我們這裡貼出核心的推理程式碼：

YoloV5 model(inputParams, trtParams, v5Params);
  if (!model.initSession(0)) {
    std::cerr << "model init failed!\n";
    exit(0);
  }

  string data_f = FLAGS_data_f;
  if (thor::os::isfile(data_f)) {
    if (thor::os::suffix(data_f) == "mp4" ||
        thor::os::suffix(data_f) == "avi") {
      std::cout << "inference on video.\n";
      VideoCapture cap(data_f);
      if (!cap.isOpened()) {
        cout << "Error opening video stream or file" << endl;
        return -1;
      }
      cap.set(CAP_PROP_POS_FRAMES, 60);

      VideoWriter video("yolov5_result.mp4", CV_FOURCC('M', 'J', 'P', 'G'), 25,
                        Size(cap.get(3), cap.get(4)));

      while (1) {
        Mat image;
        cap >> image;

        if (image.empty()) break;

        auto bboxes = model.predOneImage(image);

        std::vector<thor::Detection> dets = toThorDetections(bboxes);
        auto res =
            thor::vis::VisualizeDetections(image, dets, CLASSES, &COLORS);
        cv::imshow("yolov5", res);
        video.write(res);

        char c = (char)waitKey(25);
        if (c == 27) break;
      }
      video.release();

    }

由此可以得到我們上面給大家展示的效果。

image.png

總結

我們對JetsonNano 2GB進行了較為全面的評測，包括它的軟體相容性以及在一些極致優化的模型下的表現速度。事實上，如果用我們的512輸入尺寸的Yolov5模型的TensorRT加速，它的最高速度可以達到25ms, 這基本上可以讓你在一個嵌入式板子上將檢測模型跑到realtime.

所以當你需要部署某個應用的時候，缺乏硬體考量？JetsonNano 2GB絕對是一個不錯的首選。

哦對了對了，這次評測忘了一個很重要的點： JetsonNano 2GB的TensorCore支援fp16推理！！ 這意味著在我們的評測速度上，JetsonNano 2GB還可以做到更快的速度！

本次教程的所有相關程式碼均為我們原創，轉載請註明出處，謝謝！

相關文章