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]。

[build-dependencies]
anyhow = "1.0"
bindgen = "0.59.2"
cmake = "0.1.48"
lazy_static = "1.4.0"

目录结构

.
|____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*。

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 = ""
);

转变为

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自然也是搞不定的。

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

转变为

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

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

添加了

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<()>
whereP: 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。