Rust小技巧 - 通过FFI编程运行tensorrt模型

104 阅读 0 评论 69 点赞

我是靠谱客的博主神勇小海豚，这篇文章主要介绍Rust小技巧 - 通过FFI编程运行tensorrt模型，现在分享给大家，希望可以做个参考。

文章目录

- 1 概述
- 2 使用说明
- - 2.1 配置说明
  - 2.2 修改c++头文件
  - 2.3 编写build.rs
  - 2.4 测试
- 参考资料

1 概述

shouxieai/tensorRT_Pro是一个文档完善，效果也很不错的tensorrt库，里面有对yolov5，yolox，unet，bert，retinaface等多个常用模型的tensorrt实现。但其中的内容基本是用c++写的，比如我目前的一个服务使用的web框架是rust的，但是又想用tensorrt来部署模型，那么就需要用到rust的FFI(Foreign Function Interface)编程了。

FFI的作用就是使得rust可以像调用rust自身的函数一样去调用c/c++的函数，这种跨语言的交互要做到正确无误，就需要有一个完善的数据结构对齐的支持。也就是说我给rust一个变量的内存地址，然后告诉rust这个变量的在c/c++中的类型，rust要能把他准确无误地变成rust中的变量。这一点，rust对c语言是有官方支持的，做的很好，但是对c++就支持的不太好。

网上对于FFI编程的中文资料也非常少，这里就举个例子，做点贡献。本文对shouxieai/tensorRT_Pro中的simple-yolo做了FFI编程。

完整的代码可见https://github.com/zjuPeco/simple-yolo-sys。

2 使用说明

FFI的最终目标是写一个binding.rs文件，这个文件相当于把c/c++的头文件给翻译成了rust的头文件。如果自己手动去写这个文件的话，不是不可以，而是太麻烦，对于新接触这块的新手而言更是难上加难，网上的相关资料实在是太少了。好在有一个叫做bindgen的rust库，可以帮我们做这层转换，我们要做的就是把要转换的头文件和内容告诉它就行了。

2.1 配置说明

bindgen会去找一个叫做build.rs的文件，我们把这个文件新建在和src同级的地方。这个文件是我们告诉bindgen要做哪些转换的地方。

同时在cargo.toml当中配置好[build-dependencies]。

复制代码

1
2
3
4
5
6
[build-dependencies]
anyhow = "1.0"
bindgen = "0.59.2"
cmake = "0.1.48"
lazy_static = "1.4.0"

目录结构

复制代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
.
|____src
| |____lib.rs
|____libyolo
| |____CMakeLists.txt
| |____src
| | |____simple_yolo.cu
| | |____simple_yolo.hpp
|____build.rs
|____Cargo.toml
|____Cargo.lock
|____.gitignore
|____README.md

2.2 修改c++头文件

c语言的基本都是可以转换的，但是c++相关的就有很多转出来会有问题，一个万能的方案就是改写一下c++的代码，把一些高级的数据类型转化成c的数据类型就可以了。

我会把我对simple_yolo.hpp做的改动都在此说明一下。

使用bindgen去转换const string&是会出问题的，要全部改成const char*。

复制代码

1
2
3
4
5
6
7
8
9
10
bool compile(
    Mode mode, Type type,
    unsigned int max_batch_size,
    const string& source_onnx,
    const string& saveto,
    size_t max_workspace_size = 1<<30,
    const std::string& int8_images_folder = "",
    const std::string& int8_entropy_calibrator_cache_file = ""
);

转变为

复制代码

1
2
3
4
5
6
7
8
9
10
bool compile(
    Mode mode, Type type,
    unsigned int max_batch_size,
    const char* source_onnx,
    const char* saveto,
    size_t max_workspace_size = 1<<30,
    const char* int8_images_folder = "",
    const char* int8_entropy_calibrator_cache_file = ""
);

shared_ptr这种高级指针bindgen自然也是搞不定的。

复制代码

1
2
shared_ptr<Infer> create_infer(const string& engine_file, Type type, int gpuid, float confidence_threshold=0.25f, float nms_threshold=0.5f);

转变为

复制代码

1
2
Infer* create_infer(const char* engine_file, Type type, int gpuid, float confidence_threshold=0.25f, float nms_threshold=0.5f);

predict的左右过程都包进一个函数当中，并返回一个数据类型简单的结构体。

添加了

复制代码

1
2
3
4
5
6
7
8
9
10
11
struct Prediction{
    float** results;
    int length;
    Prediction(float** results, int length){
        this->results = results;
        this->length = length;
    }
};

Prediction* predict(Infer* engine, const cv::Mat& image);

这里的results本来想用vector的，但是使用vector也会出一些问题，某些返回来的数据会有异常，还是用二维指针吧。

如果想要干其他事情，也可以写一个数据类型只有c的简单函数，然后暴露给rust。

2.3 编写build.rs

最后把需要在rust中使用的函数在build.rs中告诉bindgen，执行cargo build就完事儿了。

复制代码

fn gen_bindings<P>(include_path: P) -> Result<()>
where
    P: AsRef<Path>,
{
    bindgen::Builder::default()
        .header(
            include_path
                .as_ref()
                .join("simple_yolo.hpp")
                .to_str()
                .ok_or_else(|| format_err!("cannot create path to darknet.hpp"))?,
        )
        .clang_arg("-I/nfs/users/chenquan/packages/tensorrt_pro/data/lean/opencv-4.2.0/include/opencv4/")
        .allowlist_function("SimpleYolo::compile")
        .allowlist_function("SimpleYolo::show_boxes")
        .allowlist_function("SimpleYolo::show_mat_shape")
        .allowlist_function("SimpleYolo::create_infer")
        .allowlist_function("SimpleYolo::predict")
        .allowlist_function("SimpleYolo::reset_engine")
        .allowlist_type("SimpleYolo::Box")
        .allowlist_type("SimpleYolo::Prediction")
        .generate()
        .map_err(|_| format_err!("failed to generate bindings"))?
        .write_to_file(&*BINDINGS_TARGET_PATH)?;
    Ok(())
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
fn gen_bindings<P>(include_path: P) -> Result<()>
where
    P: AsRef<Path>,
{
    bindgen::Builder::default()
        .header(
            include_path
                .as_ref()
                .join("simple_yolo.hpp")
                .to_str()
                .ok_or_else(|| format_err!("cannot create path to darknet.hpp"))?,
        )
        .clang_arg("-I/nfs/users/chenquan/packages/tensorrt_pro/data/lean/opencv-4.2.0/include/opencv4/")
        .allowlist_function("SimpleYolo::compile")
        .allowlist_function("SimpleYolo::show_boxes")
        .allowlist_function("SimpleYolo::show_mat_shape")
        .allowlist_function("SimpleYolo::create_infer")
        .allowlist_function("SimpleYolo::predict")
        .allowlist_function("SimpleYolo::reset_engine")
        .allowlist_type("SimpleYolo::Box")
        .allowlist_type("SimpleYolo::Prediction")
        .generate()
        .map_err(|_| format_err!("failed to generate bindings"))?
        .write_to_file(&*BINDINGS_TARGET_PATH)?;
    Ok(())
}

allowlist_function和allowlist_type表示需要进行转换的函数和数据类型。

cargo build成功之后，会在target/debug/build/simple-yolo-sys-xxx/out下生成一个bindings.rs，在lib.rs中，直接include!(concat!(env!("OUT_DIR"), "/bindings.rs"));就可以使用其中的函数了。

执行c++文件的编译也是在build.rs当中，是通过cmake这个rust库来完成的。

c++编译好的动态库在target/debug/build/simple-yolo-sys-xxx/out/lib/libcucodes.so。

2.4 测试

在lib.rs中的test_compile_tensorrt_engine可以测试onnx转tensorrt。

在lib.rs中的test_run_engine可以测试tensorrt模型能否正常使用。

使用时需要将其中的路径改成自己的路径。

测试通过就可以在其他项目中调用这个库了。

参考资料

[1] https://github.com/shouxieai/tensorRT_Pro
[2] https://github.com/alianse777/darknet-sys-rust

最后

以上就是神勇小海豚最近收集整理的关于Rust小技巧 - 通过FFI编程运行tensorrt模型的全部内容，更多相关Rust小技巧内容请搜索靠谱客的其他文章。

本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。

本文分类：rust小技巧
浏览次数：104 次浏览
发布日期：2024-09-23 17:05:01
本文链接：https://www.kaopuke.com/article/k-p-k_13_u_7_o_10_fx_12__7_k5.html

Rust小技巧 - 通过FFI编程运行tensorrt模型

文章目录

1 概述

2 使用说明

2.1 配置说明

2.2 修改c++头文件

2.3 编写build.rs

2.4 测试

参考资料

最后

评论列表共有 0 条评论

发表评论取消回复

Rust小技巧 - 通过FFI编程运行tensorrt模型

文章目录

1 概述

2 使用说明

2.1 配置说明

2.2 修改c++头文件

2.3 编写build.rs

2.4 测试

参考资料

最后

相关文章

评论列表共有 0 条评论

发表评论 取消回复

微信扫一扫：分享

发表评论取消回复