autoware.universe源码略读3.4--perception:tensorrt_yolox
Overview
可以看到,其实在最新版本的autoware中,已经把tensorrt_yolo这个包去掉了,留下的就是tensorrt_yolox,也就是用到了旷视提出的YOLOx模型的目标检测模型。而且相比于tensorrt_yolo冗余的结构,tensorrt_yolox的文件结构更加简练,没有lib文件夹,也没有那些自己定义的插件类。这里只有tensorrt_yolox.cpp一个文件就完成了主要的功能,这里来简单看下YOLOx在auto ware.universe中是如何被部署的
PS:这里的功能包似乎并不完整,看上去都没有launch文件,json文件这些,而且最新版还是比这里复杂一些的,这里还是先简单看下是怎么用的吧
结构体预定义
这个文件中先预定义了两个后边会用到的结构体,分别是Object
和GridAndStride
,这两个一个对应的是检测到的物体类,另一个对应的是YOLOx在划分网格时候会用到的结构体,对应的是网格的坐标和步长(也就是大小)
TrtYoloX
首先来看构造函数,输入的参数没有什么特殊的
const std::string & model_path // 模型文件的路径
const std::string & precision // 精度(int8,fp32这种)
const int num_class // 检测类别的数量
const float score_threshold // 设定的分数阈值
const float nms_threshold // nms的阈值
[[maybe_unused]] const std::string & cache_dir // 缓存文件的路径
const tensorrt_common::BatchConfig & batch_config // 批处理的配置
const size_t max_workspace_size // 最大空间设置
这里相比于tensorrt_yolo让我感到舒服的第一点是,终于用到了之前在common包里定义的tensorrt_common这个类,可以参考autoware.universe源码略读(1)–common。可以看到针对不同的trt_common_->getNbBindings()
函数返回结果,这里是进行了不同的处理。涉及到了Binding Indices绑定索引这个概念,在 TensorRT 中,绑定索引用于标识模型的输入和输出张量。在 YOLO 模型中,输出的预测框通常需要进行处理(例如解码、过滤等)以得到最终的检测结果。如果模型已经包含了EfficientNMS_TRT
模块,也就是对应返回值是5的情况,这些处理步骤已经在推理过程中完成,不需要额外解码。
接下来在构造函数里还进行了的一个步骤就是设置可能用到的内存空间。这里还是第一次见到std::accumulate
的这个用法,之前以为这个函数只是单纯的累加,原来最后一个参数是可以指定运算规则的,像这里指定的就是乘法了,具体的就是计算了输入的大小,+1可能是因为第一个通常是batch_size?但其实在神经网络中每次输入的肯定是一个对象
const auto input_size =
std::accumulate(input_dims.d + 1, input_dims.d + input_dims.nbDims, 1, std::multiplies<int>());
preprocess
这里是对输入图像的预处理过程,有对图像进行尺度变换的过程,涉及到的原理在之前的文章autoware.universe源码略读(3.1)–perception:yolo初识提到过,就是都宽和高都算一个比例,然后选小的比例,剩下的部分用灰色 {114, 114, 114} 给填充起来,之后用到了cv::dnn::blobFromImages
。对图像的格式细节调整了一下,可以参考OpenCV官方文档
doInference
可以看到,这里是执行具体的推理的函数。流程很简单,首先判断TrnsorRT准没准备好
if (!trt_common_->isInitialized()) {
return false;
}
然后对图像进行预处理
preprocess(images);
之后核心的部分还是被封装成了别的函数调用,当然这里还是根据之前的分类用,分类标注 就是包不包含EfficientNMS_TRT
这个模块
if (needs_output_decode_) {
return feedforwardAndDecode(images, objects);
} else {
return feedforward(images, objects);
}
feedforward
这个是包含EfficientNMS_TRT
模块时的主要步骤的函数,其实核心调用的是一句话
trt_common_->enqueueV2(buffers.data(), *stream_, nullptr);
如果有印象的话,这里调用的是nvinfer1::IExecutionContext
对象的enqueueV2
函数,所以其实还是单纯地封装了几层,剩下的就是对输出的处理了,把输出的结果转换成object
for (size_t i = 0; i < batch_size; ++i) {
const size_t num_detection = static_cast<size_t>(out_num_detections[i]);
ObjectArray object_array(num_detection);
for (size_t j = 0; j < num_detection; ++j) {
Object object{};
const auto x1 = out_boxes[i * max_detections_ * 4 + j * 4] / scales_[i];
const auto y1 = out_boxes[i * max_detections_ * 4 + j * 4 + 1] / scales_[i];
const auto x2 = out_boxes[i * max_detections_ * 4 + j * 4 + 2] / scales_[i];
const auto y2 = out_boxes[i * max_detections_ * 4 + j * 4 + 3] / scales_[i];
object.x_offset = std::clamp(0, static_cast<int32_t>(x1), images[i].cols);
object.y_offset = std::clamp(0, static_cast<int32_t>(y1), images[i].rows);
object.width = static_cast<int32_t>(std::max(0.0F, x2 - x1));
object.height = static_cast<int32_t>(std::max(0.0F, y2 - y1));
object.score = out_scores[i * max_detections_ + j];
object.type = out_classes[i * max_detections_ + j];
object_array.emplace_back(object);
}
objects.emplace_back(object_array);
}
feedforwardAndDecode
这个函数对应的就是没有包含EfficientNMS_TRT
模块时的主要步骤,推理的核心步骤是一样的,执行的是
trt_common_->enqueueV2(buffers.data(), *stream_, nullptr);
主要的差别体现在对输出的处理上,因为这里的输出应该是没有处理过的,所以后边是使用了decodeOutputs
对输出的结果又进行了一步解码的处理
for (size_t i = 0; i < batch_size; ++i) {
auto image_size = images[i].size();
float * batch_prob = out_prob_h_.get() + (i * out_elem_num_per_batch_);
ObjectArray object_array;
decodeOutputs(batch_prob, object_array, scales_[i], image_size);
objects.emplace_back(object_array);
}
decodeOutputs
解码处理这里涉及到的内容还挺多的,这里首先对应的两个步骤就是生成网格和候选框
generateGridsAndStride(input_width, input_height, strides, grid_strides);
generateYoloxProposals(grid_strides, prob, score_threshold_, proposals);
之后会对候选框根据置信度进行排序
qsortDescentInplace(proposals);
然后再执行一下nms非极大值抑制的步骤
nmsSortedBboxes(proposals, picked, nms_threshold_);
最后也是把检测结果赋值给objects
对象
for (int i = 0; i < count; i++) {
objects[i] = proposals[picked[i]];
// adjust offset to original unpadded
float x0 = (objects[i].x_offset) / scale;
float y0 = (objects[i].y_offset) / scale;
float x1 = (objects[i].x_offset + objects[i].width) / scale;
float y1 = (objects[i].y_offset + objects[i].height) / scale;
// clip
x0 = std::clamp(x0, 0.f, static_cast<float>(img_size.width - 1));
y0 = std::clamp(y0, 0.f, static_cast<float>(img_size.height - 1));
x1 = std::clamp(x1, 0.f, static_cast<float>(img_size.width - 1));
y1 = std::clamp(y1, 0.f, static_cast<float>(img_size.height - 1));
objects[i].x_offset = x0;
objects[i].y_offset = y0;
objects[i].width = x1 - x0;
objects[i].height = y1 - y0;
}
这里来看下generateYoloxProposals
这个函数,这里在解码候选框的时候用到了指数运算,这是因为YOLO算法中,边界框的宽度和高度的计算采用的是指数映射,所以这里的代码是这样的
float x_center = (feat_blob[basic_pos + 0] + grid0) * stride;
float y_center = (feat_blob[basic_pos + 1] + grid1) * stride;
float w = exp(feat_blob[basic_pos + 2]) * stride;
float h = exp(feat_blob[basic_pos + 3]) * stride;
float x0 = x_center - w * 0.5f;
float y0 = y_center - h * 0.5f;
这里对分数降序排列的算法在qsortDescentInplace
里,因为自己数据结构的知识相对比较薄弱,正好这里看一下
首先是规定了左右边界和中位值(索引在一半位置的值)
int i = left;
int j = right;
float p = faceobjects[(left + right) / 2].score;
之后进行第一次划分的目的是:把数值根据相对中位元素的大小分别放在两边
while (i <= j) {
while (faceobjects[i].score > p) {
i++;
}
while (faceobjects[j].score < p) {
j--;
}
if (i <= j) { // 把大的放在一边,小的也放在一边
// swap
std::swap(faceobjects[i], faceobjects[j]);
i++;
j--;
}
}
之后进行递归调用,相当于两边分开进行递归排序了
#pragma omp parallel sections
{
#pragma omp section
{
if (left < j) {
qsortDescentInplace(faceobjects, left, j);
}
}
#pragma omp section
{
if (i < right) {
qsortDescentInplace(faceobjects, i, right);
}
}
}
最后nms这里,这里只是简单计算了下IoU,然后根据IoU和阈值判断这个是不是应该keep
,似乎是没有涉及到类别概念之类的
tensorrt_yolox_node
galatic版本下,这个节点文件是有定义的,但是我看了下是没有使用的示例。这里的构造函数没有什么特别的,首先是加载参数,包括模型文件路径、标签文件路径、精度、分数阈值和nms阈值,所以如果自己写launch的话把这几个参数设置一下应该就行了
接下来就是对象的实例化以及话题订阅发布之类的了,这个和tensorrt_yolo的类似,图像话题的订阅也是通过一个timer_
来控制的,这样能确保只有有人订阅的时候再去订阅图像,这样能节省一定的资源。这里定额与时用到的回调函数是onImage
,调用doInference
来执行推理过程
if (!trt_yolox_->doInference({in_image_ptr->image}, objects)) {
RCLCPP_WARN(this->get_logger(), "Fail to inference");
return;
}
下面也是根据检测到的类别进行图像上标记以及话题发布
yolox_single_image_inference_node
这个看起来是针对一张图片的时候的一个检测节点,在构造函数里,少了标签文件这个参数,多了一个是否保存生成图片的参数,这里自己写一个launch来试一下,launch文件很好写,只需要在新建一个launch文件,然后在里面创建一个 .launch.xml 后缀的文件就好了
<launch>
<node pkg="tensorrt_yolox" exec="yolox_single_image_inference" name="$(anon tensorrt_yolox)" output="screen">
<param name="image_path" value="path_to_image"/>
<param name="model_path" value="$(find-pkg-share tensorrt_yolox)/data/yolox-tiny.onnx"/>
<param name="precision" value="FP32"/>
<param name="save_image" value="true"/>
</node>
</launch>
然后在cmake里记得修改下面这句话
ament_auto_package(INSTALL_TO_SHARE
data
launch
)
这样才能找到launch文件,之后只需要重新编译一下这个包就好了
colcon build --symlink-install --cmake-args -DCMAKE_BUILD_TYPE=Release --packages-select tensorrt_yolox
然后再执行命令
ros2 launch tensorrt_yolox single_yolox.launch.xml
OK会有很多输出,不过我感觉这个运行速度也不是很快啊?
我这里扔进去的就是YOLO中那张经典的狗的照片,输出出来的效果其实一般,自行车的预测框明显位置偏了些,可能是占比太大了?
总结
tensorrt_yolox和tensorrt_yolo的部署其实看下来都不算复杂,主要涉及到的可能是利用TensorRT这个库,其实感觉不懂YOLO原理也一样可以部署。自己也是把yolox的这个示例试了一下,不过感觉速度和精度都没有自己想想的好,摩托车和后边的车还好,主要是前面的这个自行车,可能还有哪里自己没有考虑到吧,后续自己使用的时候再留意一下
以及,最新版本中的tensorrt_yolox模块似乎更加合理,如果后边有确定要用的话,或许可以看看把最新版本的移植过来?只要注意下其他的接口就好了