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

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