[YOLOv8] 在Jetson Nano当中部署YOLOv8目标检测算法

梯度下降法_

已于 2024-04-24 20:18:27 修改

阅读量3.2k

点赞数 16

文章标签：计算机视觉深度学习 YOLO 嵌入式硬件

于 2024-04-24 20:16:18 首次发布

本文链接：https://blog.csdn.net/m0_75272311/article/details/137983073

版权

源码下载

本文所有源码在GitHub当中发布，喜欢的话麻烦点个star

GitHub - Dominic23331/yolov8_tensorrt

前期准备

硬件准备

在前期的准备当中，你需要有以下硬件：

个人PC
Jetson Nano开发板
TF卡
路由器
网线

软件准备

在个人电脑当中，你需要安装以下软件：

接下来主要讲解一下jetson nano当中的软件安装过程。

1、首先，你需要准备一个 Jetson Nano（B01或A02 均可，建议采用 4GB 版的开发板），以及一张容量大于 16GB 的 TF 卡。

2、在英伟达网站中下载jetson nano的镜像文件。

3、使用balenaEtcher软件将镜像写入TF卡当中。

4、将TF卡插入jetson nano当中，并进行初始化设置。

5、使用网线或无线连接将 Jetson Nano 与路由器相连接，然后在浏览器中输入路由器的管理界面地址（通常是 192.168.0.1）。从管理界面中查找 Jetson Nano 的 IP 地址。

6、使用 MobaXterm 软件的 SSH 功能，输入 Jetson Nano 的 IP 地址以及账号密码，连接到 Jetson Nano。

7、（可选）为了方便之后的模型转换操作，建议将trtexec加入环境变量中。

sudo gedit ~/.bashrc

# 在.bashrc当中添加以下内容

export PATH=/usr/src/tensorrt/bin:$PATH

# 更新环境变量

source ~/.bashrc

远程开发的准备

为了方便开发，我们将使用 Clion 配置远程开发，这样就可以在个人电脑上编写 Jetson Nano 的代码。在配置过程中，请确保 Jetson Nano 能够稳定连接到路由器，建议使用网线连接以提高稳定性。以下是配置步骤：

1、打开Clion，并创建一个工程，在菜单栏中选择setting，并打开build,Execution,Deployment->Toolchains。在Toolchains当中新建一个Remote Host，通过指引配置jetson nano的远程开发环境。

2、点击上图红圈位置以打开远程服务器的配置栏，然后输入 Jetson Nano 的 IP 地址和账号密码来建立连接，这样就可以在 Clion 中进行对 Jetson Nano 的远程开发了。

模型的转换

pt转换为onnx

为了将模型部署到jetson nano当中，我们首先需要将需要转换的模型导出为onnx格式。首先，你需要下载YOLOv8的模型文件，模型可以从这里下载。由于jetson nano的GPU计算能力较弱，在这里我使用了YOLOv8n模型，并将输入图像的尺寸缩小为原来的四分之一。转换的代码如下所示：

from ultralytics import YOLO

model = YOLO("yolov8n.pt")
model.export(imgsz=320, format='onnx')

这样，我们就得到了onnx格式的YOLOv8模型了。

onnx转换为engine

在 Jetson Nano 上，我们可以使用 TensorRT 对模型进行加速并部署。由于 TensorRT 对模型的优化与硬件有关，因此需要将 ONNX 模型上传至 Jetson Nano，并通过 trtexec 工具进行模型的转换。以下是模型转换的命令：

trtexec --onnx=onnx模型路径 --saveEngine=模型保存路径

通过以上操作，即可将YOLOv8模型转换为jetson nano支持的格式了。

模型的部署

前期准备

在模型部署当中，首先我们需要将模型文件读取至内存当中，并初始化TensorRT模型，为了方便编程，本文将YOLOv8模型封装为一个类当中，通过构造函数的方式对模型进行初始化操作。

YOLO::YOLO(const std::string &model_path, nvinfer1::ILogger &logger) {
    // load engine file
    std::ifstream engineStream(model_path, std::ios::binary);

    // check engine file exist.
    if (!engineStream.is_open()) {
        std::cout << "Cannot find model from: " + model_path << std::endl;
        exit(0);
    }

    engineStream.seekg(0, std::ios::end);
    const size_t modelSize = engineStream.tellg();
    engineStream.seekg(0, std::ios::beg);
    std::unique_ptr<char[]> engineData(new char[modelSize]);
    engineStream.read(engineData.get(), modelSize);
    engineStream.close();

    runtime = nvinfer1::createInferRuntime(logger);
    engine = runtime->deserializeCudaEngine(engineData.get(), modelSize);
    context = engine->createExecutionContext();

    context->setBindingDimensions(0, nvinfer1::Dims4(1, 3, input_h, input_w));

    cudaStreamCreate(&stream);

    offset[0] = 0;
    offset[1] = 0;

    // Calculate output dimension by inputting image size
    out_dim_2 = (input_w / 8) * (input_h / 8) + (input_w / 16) * (input_h / 16) + (input_w / 32) * (input_h / 32);
    boxes_result.resize(out_dim_2 * 84);
}

为了展示加载的模型的输入输出，我们编写了一个函数，用于打印模型的输入和输出信息。

void YOLO::show() {
    for (int i = 0; i < engine->getNbBindings(); i++) {
        std::cout << "node: " << engine->getBindingName(i) << ", ";
        if (engine->bindingIsInput(i)) {
            std::cout << "type: input" << ", ";
        } else {
            std::cout << "type: output" << ", ";
        }
        nvinfer1::Dims dim = engine->getBindingDimensions(i);
        std::cout << "dimensions: ";
        for (int d = 0; d < dim.nbDims; d++) {
            std::cout << dim.d[d] << " ";
        }
        std::cout << "\n";
    }
}

/**输出结果
node: images, type: input, dimensions: 1 3 320 320 
node: output0, type: output, dimensions: 1 84 2100 
*/

接下来，我们编写一个预热函数，通过喂入模型随机数据，使模型在GPU当中预热，加快模型的推理速度。

void YOLO::warmup(int epoch)
{
    std::cout << "Warm up." << std::endl;

    for (int step = 0; step <= epoch; ++step)
    {
        // use a random tensor to warmup
        cv::Mat randomImage(input_h, input_w, CV_8UC3);
        cv::randu(randomImage, cv::Scalar::all(0), cv::Scalar::all(255));
        run(randomImage);

        printProgressBar(step, epoch);
    }
    std::cout << std::endl;
}

通过以上代码，经过多轮预热过后，模型能够顺畅的在GPU当中计算。

模型推理

接下来，我们编写模型的推理代码，推理代码主要由预处理，模型推理以及后处理所构成，在预处理阶段，代码首先会将输入的图像进行resize，将图像调整为320x320大小。由于OpenCV在读取图像时，会以BGR的顺序将图像读入内存，因此在输入到模型前，应将图像从BGR格式转换为RGB的格式，并将每一个像素除以255。以下是预处理代码：

std::vector<float> YOLO::preprocess(cv::Mat &image)
{
    // get resized image
    std::tuple<cv::Mat, int, int> resized = resize(image, input_w, input_h);
    cv::Mat resized_image = std::get<0>(resized);

    // get resize offset
    offset[0] = std::get<1>(resized);
    offset[1] = std::get<2>(resized);

    // BGR2RGB
    cv::cvtColor(resized_image, resized_image, cv::COLOR_BGR2RGB);

    std::vector<float> input_tensor;
    for (int k = 0; k < 3; k++) {
        for (int i = 0; i < resized_image.rows; i++) {
            for (int j = 0; j < resized_image.cols; j++) {
                input_tensor.emplace_back(((float) resized_image.at<cv::Vec3b>(i, j)[k]) / 255.);
            }
        }
    }

    return input_tensor;
}

在模型的推理阶段，我们所要做的是将预处理的图像从CPU当中转到GPU中，并将图像输入到模型当中，经过模型的推理之后，再将推理的结果从GPU搬运到CPU当中进行后处理的计算。以下代码为推理程序：

// upload to cuda
    cudaMalloc(&buffer[0], 3 * input_h * input_w * sizeof(float));
    cudaMalloc(&buffer[1], out_dim_2 * 84 * sizeof(float));
    cudaMemcpyAsync(buffer[0], input.data(), 3 * input_h * input_w * sizeof(float), cudaMemcpyHostToDevice, stream);

    // inference
    context->enqueueV2(buffer, stream, nullptr);
    cudaStreamSynchronize(stream);

    // download from cuda
    cudaMemcpyAsync(boxes_result.data(), buffer[1], out_dim_2 * 84 * sizeof(float), cudaMemcpyDeviceToHost);

经过了模型的推理，我们可以获得模型输出的所有的检测框，由于模型输出的维度为（1，84，2100），而TensorRT输出的结果为一个一维向量，因此需要通过两个for循环，对模型的输出进行解码。

std::vector<float> b;
    for (int i=0; i < out_dim_2; i++)
    {
        b.clear();
        for (int j=0; j < 84; j++)
        {
            b.push_back(tensor[j * out_dim_2 + i]);
        }

        b = decode_cls(b);
        if (b[4] < conf_threshold)
            continue;

        Box db;
        db.x1 = b[0] - b[2] / 2;
        db.y1 = b[1] - b[3] / 2;
        db.x2 = b[0] + b[2] / 2;
        db.y2 = b[1] + b[3] / 2;
        db.conf = b[4];
        db.cls = classes[(int) b[5]];

        boxes.push_back(db);
    }

我们通过以上代码，即可得到大于置信度阈值的所有检测框。由于检测框当中存在重叠，因此需要通过非极大值抑制算法，对冗余的检测框进行过滤。以下是算法的代码：

/**
 * decode classes
 * @param box output box
 * @return boxes
 */
std::vector<float> decode_cls(std::vector<float>& box)
{
    std::vector<float> cls_list(box.begin() + 4, box.end());
    float conf = *std::max_element(cls_list.begin(), cls_list.end());
    float cls = std::max_element(cls_list.begin(), cls_list.end()) - cls_list.begin();

    std::vector<float> result(box.begin(), box.begin() + 4);
    result.push_back(conf);
    result.push_back(cls);

    return result;
}


/**
 * compare two boxes
 * @param b1 box1
 * @param b2 box2
 * @return result
 */
bool compare_boxes(const Box& b1, const Box& b2)
{
    return b1.conf < b2.conf;
}


/**
 * compare two boxes IOU
 * @param b1 box1
 * @param b2 box2
 * @return iou
 */
float intersection_over_union(const Box& b1, const Box& b2)
{
    float x1 = std::max(b1.x1, b2.x1);
    float y1 = std::max(b1.y1, b2.y1);
    float x2 = std::min(b1.x2, b2.x2);
    float y2 = std::min(b1.y2, b2.y2);

    // get intersection
    float box_intersection = std::max((float)0, x2 - x1) * std::max((float)0, y2 - y1);

    // get union
    float area1 = (b1.x2 - b1.x1) * (b1.y2 - b1.y1);
    float area2 = (b2.x2 - b2.x1) * (b2.y2 - b2.y1);
    float box_union = area1 + area2 - box_intersection;

    // To prevent the denominator from being zero, add a very small numerical value to the denominator
    float iou = box_intersection / (box_union + 0.0001);

    return iou;
}


/**
 * NMS
 * @param boxes output boxes from YOLO
 * @param iou_thre  iou threshold
 * @return boxes after NMS
 */
std::vector<Box> non_maximum_suppression(std::vector<Box> boxes, float iou_thre)
{
    // Sort boxes based on confidence
    std::sort(boxes.begin(), boxes.end(), compare_boxes);

    std::vector<Box> result;
    std::vector<Box> temp;
    while (!boxes.empty())
    {
        temp.clear();

        Box chosen_box = boxes.back();
        boxes.pop_back();
        for (auto & boxe : boxes)
        {
            if (boxe.cls != chosen_box.cls || intersection_over_union(boxe, chosen_box) < iou_thre)
                temp.push_back(boxe);
        }

        boxes = temp;
        result.push_back(chosen_box);
    }
    return result;
}

最后，即可将输出的检测框还原至原图的尺寸。

// resize boxes
for (auto & boxe : boxes)
{
    boxe.x1 = MAX((boxe.x1 - offset[0]) * img_w / (input_w - 2 * offset[0]), 0);
    boxe.y1 = MAX((boxe.y1 - offset[1]) * img_h / (input_h - 2 * offset[1]), 0);
    boxe.x2 = MIN((boxe.x2 - offset[0]) * img_w / (input_w - 2 * offset[0]), img_w);
    boxe.y2 = MIN((boxe.y2 - offset[1]) * img_h / (input_h - 2 * offset[1]), img_h);
}

我们将图像输入模型，即可得到结果：

总结

经过上述步骤，我们就完成了在jetson nano当中部署YOLOv8的工作。为了检验其效率，我编写了代码对每个步骤的耗时进行统计，以下是统计结果：

Waste time: 0.07
Preprocess time: 0.03
Cuda upload time: 0.00
Model inference time: 0.01
Cuda download time: 0.00
Postprocess time: 0.03
FPS: 14.38

由结果可见，推理一张图像的总耗时大概70ms左右，cuda搬运数据几乎不耗时，大概是因为jetson nano的GPU和CPU共用一个内存吧。而令人困惑的是，模型的推理时间只有10ms，反而前后处理消耗了60ms的时间，可能是包含了多个大循环所导致的，接下来将尝试使用CUDA对前后处理进行加速。

梯度下降法_

关注

16
点赞
踩
58

收藏

觉得还不错? 一键收藏
1
评论
[YOLOv8] 在Jetson Nano当中部署YOLOv8目标检测算法

使用jetson nano部署YOLOv8
复制链接

扫一扫