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<()>
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