写在前面
本篇文章程序链接：
上一节主要讲了模型的创建、加载与推理，本节主要是各个模型的预处理与后处理部分。
分类网络
模型训练导出
假设现在有个分类网络，如
torchvision自带的预训练模型，模型具体结构：class Classifier(torch.nn.Module): def __init__(self): super().__init__() self.backbone = torchvision.models.resnet18(pretrained=True) def forward(self, x): feature = self.backbone(x) probability = torch.softmax(feature, dim=1) return probability dummy = torch.zeros(1, 3, 224, 224) torch.onnx.export(model, (dummy,), "workspace/classifier.onnx", input_names=["image"], output_names=["prob"], dynamic_axes={"image": {0: "batch"}, "prob": {0: "batch"}}, opset_version=11)模型编译与加载，onnx模型的编译与加载在第三节已详细讲解，本节不再过多赘述。
数据预处理
网络输入主要过程：加载图像 —> 缩放 —> 图像标准化 —> cpu to gpu，具体过程封装成一个方法，适用于一般预处理流程：
std::pair<float *, float *> preprocess(int b, int c, int w, int h, const cv::Mat &img, cudaStream_t stream) { if (b <= 0 || c <= 0 || w <= 0 || h <= 0 || img.empty()) { return std::make_pair(nullptr, nullptr); } cv::Mat _img(h, w, CV_8UC3); // 图像缩放 if (img.rows != h || img.cols != w) { cv::resize(img, _img, cv::Size(w, h)); } else { img.copyTo(_img); } // 初始化内存地址，分配内存 int num = b * c * w * h; float *in_data_host = nullptr; float *in_data_device = nullptr; cudaMallocHost(&in_data_host, sizeof(float) * num); cudaMalloc(&in_data_device, sizeof(float) * num); // 图像标准化，bgr to rgb，并且将img data 拷贝到开辟的内存上 int _img_size = w * h; float mean[] = {0.406, 0.456, 0.485}; float std[] = {0.225, 0.224, 0.229}; unsigned char *pimg = _img.data; float *_b = in_data_host + _img_size * 0; float *_g = in_data_host + _img_size * 1; float *_r = in_data_host + _img_size * 2; for (int i = 0; i < _img_size; ++i, pimg += 3) { *_r++ = (pimg[0] / 255.0f - mean[0]) / std[0]; *_g++ = (pimg[1] / 255.0f - mean[1]) / std[1]; *_b++ = (pimg[2] / 255.0f - mean[2]) / std[2]; } // cpu to gpu if (stream == nullptr) { cudaMemcpyAsync(in_data_device, in_data_host, num * sizeof(float), cudaMemcpyHostToDevice); } else { cudaMemcpyAsync(in_data_device, in_data_host, num * sizeof(float), cudaMemcpyHostToDevice, stream); } return std::make_pair(in_data_host, in_data_device); }数据后处理
float *prob = output_data_host; int predict_label = std::max_element(prob, prob + num_classes) - prob; // 确定预测分数最大值的下标 auto predict_name = labels[predict_label]; // labels 为解析后的类别数组
Yolov5推理
数据预处理。预处理部分主要包括图像的缩放与归一化；
std::pair<cv::Mat, cv::Mat> letterbox(const cv::Mat &img, const cv::Size new_shape, const cv::Scalar color, float *d2i) { // Scale ratio float scale = std::min((new_shape.width / (float)img.cols), (new_shape.height / (float)img.rows)); float i2d[6]; /* * [scale, 0] 控制的是缩放 * [0, scale] * i2d[2] 与 i2d[5] 控制的是两次平移 */ i2d[0] = scale; i2d[1] = 0; i2d[2] = (-scale * img.cols + new_shape.width + scale - 1) * 0.5; i2d[3] = 0; i2d[4] = scale; i2d[5] = (-scale * img.rows + new_shape.height + scale - 1) * 0.5; cv::Mat m2x3_i2d(2, 3, CV_32F, i2d); cv::Mat m2x3_d2i(2, 3, CV_32F, d2i); // 反仿射变换，input_image to img 用的 cv::invertAffineTransform(m2x3_i2d, m2x3_d2i); cv::Mat input_image(new_shape.height, new_shape.width, CV_8UC3); cv::warpAffine(img, input_image, m2x3_i2d, input_image.size(), cv::INTER_LINEAR, cv::BORDER_CONSTANT, color); // 此处返回 m2x3_d2i 是为了方便分割模型推理完后，将分割结果还原为原图尺寸 return std::make_pair(input_image, m2x3_d2i); } std::tuple<float *, float *> preprocess(const std::tuple<int, int, int, int> input_shape, const cv::Mat &image, float *d2i) { auto [b, c, w, h] = input_shape; int i_num = b * c * w * h; float *input_data_host = nullptr; float *input_data_device = nullptr; CTS::checkRuntime(cudaMallocHost(&input_data_host, i_num * sizeof(float))); CTS::checkRuntime(cudaMalloc(&input_data_device, i_num * sizeof(float))); auto res = CTS::letterbox(image, cv::Size(w, h), cv::Scalar(114, 114, 114), d2i); cv::Mat input_image = res.first; int img_size = input_image.cols * input_image.rows; unsigned char *pimg = input_image.data; float *phost_b = input_data_host + img_size * 0; float *phost_g = input_data_host + img_size * 1; float *phost_r = input_data_host + img_size * 2; for (int i = 0; i < img_size; ++i, pimg += 3) { *phost_r++ = pimg[0] / 255.0f; *phost_g++ = pimg[1] / 255.0f; *phost_b++ = pimg[2] / 255.0f; /* YOLOX 的处理过程 * *phost_r++ = pimg[0]; * *phost_g++ = pimg[1]; * *phost_b++ = pimg[2]; */ } return std::make_tuple(input_data_host, input_data_device); }结果后处理。后处理主要包含
NMS与结果还原到原图；std::vector<std::vector<float>> nms(std::vector<std::vector<float>> bboxes, float nms_threshold) { // 大于降序，小于升序 std::sort(bboxes.begin(), bboxes.end(), [](std::vector<float> &a, std::vector<float> &b) { return a[5] > b[5]; }); std::vector<bool> remove_flags(bboxes.size()); std::vector<std::vector<float>> boxes_res; boxes_res.reserve(bboxes.size()); auto iou = [](const std::vector<float> &a, const std::vector<float> &b) { float cross_left = std::max(a[0], b[0]); float cross_top = std::max(a[1], b[1]); float cross_right = std::min(a[2], b[2]); float cross_bottom = std::min(a[3], b[3]); float cross_area = std::max(0.0f, cross_right - cross_left) * std::max(0.0f, cross_bottom - cross_top); float union_area = std::max(0.0f, a[2] - a[0]) * std::max(0.0f, a[3] - a[1]) + std::max(0.0f, b[2] - b[0]) * std::max(0.0f, b[3] - b[1]) - cross_area; if (cross_area == 0.0f || union_area == 0.0f) { return 0.0f; } return cross_area / union_area; }; for (int i = 1; i < bboxes.size(); ++i) { if (remove_flags[i]) { continue; } auto &ibox = bboxes[i]; boxes_res.emplace_back(ibox); for (int j = i + 1; j < bboxes.size(); ++j) { if (remove_flags[j]) { continue; } auto &jbox = bboxes[j]; if (jbox[4] == ibox[4]) { if (iou(ibox, jbox) >= nms_threshold) { remove_flags[j] = true; } } } } printf("boxes_res.size = %d\n", boxes_res.size()); return boxes_res; } std::vector<std::vector<float>> postprocess(float *output_data, int numbox, int numprob, int numcls, const float *d2i) { std::vector<std::vector<float>> bboxes; float confidence_threshold = 0.25; for (int i = 0; i < numbox; ++i) { float *ptr = output_data + i * numprob; float objness = ptr[4]; if (objness < confidence_threshold) { continue; } float *pclass = ptr + 5; int label = std::max_element(pclass, pclass + numcls) - pclass; float prob = pclass[label]; float confidence = prob * objness; if (confidence < confidence_threshold) { continue; } // xywh float cx = ptr[0]; float cy = ptr[1]; float w = ptr[2]; float h = ptr[3]; // xywh to xyxy float left = cx - w * 0.5; float top = cy - h * 0.5; float right = cx + w * 0.5; float bottom = cy + w * 0.5; // 现在的结果是在缩放后的预处理图像上，将结果转换到输入的原图上去 float base_img_left = d2i[0] * left + d2i[2]; float base_img_right = d2i[0] * right + d2i[2]; float base_img_top = d2i[0] * top + d2i[5]; float base_img_bottom = d2i[0] * bottom + d2i[5]; bboxes.push_back({base_img_left, base_img_top, base_img_right, base_img_bottom, (float)label, confidence}); } float nms_threshold = 0.5; auto boxes = nms(bboxes, nms_threshold); return boxes; }
Unet分割
以Unet官方模型为例，预处理部分与 Yolov5 预处理相同，因此只考虑后处理部分。
模型推理后处理过程为：
- 模型推理得到一个 [batch，512，512，21] 的
tensor，表示每个像素预测对应类别的概率值；假设batch = 1； - 提取每个像素最大的概率值与对应的类别id，得到2个 [512, 512] 的
tensor，然后将这两个结果缩放至原图像大小； - 根据类别id生成对应掩码图并与原图像混合。
cv::Mat postprocess(float *output_data_host, int output_width, int output_height, int num_class, int batch, const cv::Mat &img, const cv::Mat &m2x3_d2i) {
cv::Mat out_prob(output_height, output_width, CV_32F); // 存储的是每个像素点的最大概率值。
cv::Mat out_iclass(output_height, output_width, CV_8U); // 存储的是每个像素点预测的类别索引。
float *pnet = output_data_host + batch * output_width * output_height * num_class; // 指向第 batch 个输出数据的首地址。
float *prob = out_prob.ptr<float>(0);
std::uint8_t *pidx = out_iclass.ptr<std::uint8_t>(0);
std::set<int> st{};
// 提取每个像素最大的概率值与对应的类别id值
for (int k = 0; k < out_prob.cols * out_prob.rows; ++k, pnet += num_class, ++prob, ++pidx) {
int ic = std::max_element(pnet, pnet + num_class) - pnet; // 获取最大概率对应的索引
*prob = pnet[ic]; // 最大概率值
*pidx = ic; // 最大概率类被索引
}
// 将结果缩放至原图像大小
cv::warpAffine(out_prob, out_prob, m2x3_d2i, img.size(), cv::INTER_LINEAR);
cv::warpAffine(out_iclass, out_iclass, m2x3_d2i, img.size(), cv::INTER_LINEAR);
cv::Mat image = img.clone();
auto pimage = image.ptr<cv::Vec3b>(0);
auto pprob = out_prob.ptr<float>(0);
auto pclass = out_iclass.ptr<uint8_t>(0);
// 根据类别id生成对应掩码图并与原图像混合
for (int i = 0; i < img.cols * img.rows; ++i, ++pimage, ++pprob, ++pclass) {
int iclass = *pclass;
float probability = *pprob;
auto &pixel = *pimage;
float foreground = std::min(0.6f + probability * 0.2f, 0.8f);
float background = 1 - foreground;
for (int c = 0; c < 3; ++c) {
auto value = pixel[c] * background + foreground * _classes_colors[iclass * 3 + 2 - c];
pixel[c] = (int)std::min(value, 255.0f);
}
}
return image;
}AlphaPose
本次测试采用 AlphaPose 下的 Fast Pose 模型。人体姿态检测过程：模型首先检测人的位置，然后再检测人体姿态关键点。本次测试中人体位置本次不在检测，直接传入人体位置对应的box框坐标，只进行人体姿态关键点的测试。模型的输出信息为：关键点的xy坐标，关键点的置信度（x, y, confidence……）。
数据预处理。与
Yolov5预处理类似，只是加了padding处理。float rate = box.width > 100 ? 0.1f : 0.15f; float pad_width = box.width * (1 + 2 * rate); float pad_height = box.height * (1 + 2 * rate); float scale = std::min((float)w / pad_width, (float)h / pad_height);结果后处理。模型结果后处理比较复杂时，可以在导出onnx模型时，将后处理过程也添加到onnx模型中，这样不仅可以省去复杂的后处理过程，还可以提升处理速度。
导出onnx模型时添加后处理过程，只需将后处理的程序添加到导出时的forward之后即可。
import torch import yaml from easydict import EasyDict as edict from alphapose.models import builder class Alphapose(torch.nn.Module): def __init__(self): super().__init__() config_file = "configs/halpe_136/resnet/256x192_res50_lr1e-3_2x-regression.yaml" check_point = "pretrained_models/multi_domain_fast50_regression_256x192.pth" with open(config_file, "r") as f: config = edict(yaml.load(f, Loader=yaml.FullLoader)) self.pose_model = builder.build_sppe(config.MODEL, preset_cfg=config.DATA_PRESET) self.pose_model.load_state_dict(torch.load(check_point, map_location="cpu")) def forward(self, x): hm = self.pose_model(x) # 不加后处理则直接返回 hm # 后处理程序 stride = int(x.size(2) / hm.size(2)) b, c, h, w = map(int, hm.size()) prob = hm.sigmoid() confidence, _ = prob.view(-1, c, h * w).max(dim=2, keepdim=True) prob = prob / prob.sum(dim=[2, 3], keepdim=True) coordx = torch.arange(w, device=prob.device, dtype=torch.float32) coordy = torch.arange(h, device=prob.device, dtype=torch.float32) hmx = (prob.sum(dim=2) * coordx).sum(dim=2, keepdim=True) * stride hmy = (prob.sum(dim=3) * coordy).sum(dim=2, keepdim=True) * stride return torch.cat([hmx, hmy, confidence], dim=2) model = Alphapose().eval() dummy = torch.zeros(1, 3, 256, 192) torch.onnx.export( model, (dummy,), "fast-pose-136.onnx", input_names=["images"], output_names=["keypoints"], opset_version=11, dynamic_axes={ "images": {0: "batch"}, "keypoints": {0: "batch"} } )模型结果解析。
cv::Mat fast_pose_postprocess(float *output_data_host, cv::Mat &image, const int keypoint_num, const int keypoint_info, const float *d2i) { cv::Mat res = image.clone(); for (int i = 0; i < keypoint_info; ++i) { float *pkey = output_data_host + i * keypoint_num; float x = pkey[0]; // (x, y) 关键点坐标 float y = pkey[1]; float confidence = pkey[2]; // 置信度 if (confidence > 0.5) { x = x * d2i[0] + d2i[2]; y = y * d2i[0] + d2i[5]; cv::circle(res, cv::Point(x, y), 2, cv::Scalar(0, 255, 0), -1, 16); } } return res; }
mmdetection
未完，待续~
ONNXRuntime推理
本节简要描述利用 ONNX Runtime 部署 ONNX 模型过程，其具体流程为：
- 初始化 ONNX Runtime 环境：
- 创建 ONNX Runtime 的
Env对象，用于配置日志级别和会话标识。 - 使用
SessionOptions设置推理参数，如线程数、优化等级等。
- 创建 ONNX Runtime 的
- 加载模型：通过
Ort::Session创建推理会话，加载.onnx模型文件。 - 获取模型的输入输出信息：使用
GetInputTypeInfo和GetOutputTypeInfo获取输入和输出的类型信息。提取输入输出张量的维度信息，用于后续张量创建。 - 准备输入张量：
- 根据模型输入的维度创建输入张量。
- 填充输入数据（如图像预处理、归一化等）。
- 使用
Ort::Value::CreateTensor创建输入张量。
- 准备输出张量：获取输出张量的形状信息，创建相应大小的输出张量。
- 推理过程：调用
session.Run方法执行推理，将输入数据送入模型并获取输出。参数包括输入输出节点名称、输入输出张量、以及推理选项。 - 输出结果解析。
void onnxruntime_infer() {
// 初始化 ONNX Runtime 环境
Ort::Env env(ORT_LOGGING_LEVEL_INFO, "onnx");
Ort::SessionOptions session_options;
auto mem = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
session_options.SetInterOpNumThreads(1);
session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
// 加载模型
std::string model_path = "../src/cuda-tensorrt-senior-api/static/yolov5s_modified.onnx";
std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
std::wstring wide_model_path = converter.from_bytes(model_path);
Ort::Session session(env, wide_model_path.c_str(), session_options);
auto output_dims = session.GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
const char *input_names[] = {"images"}, *output_names[] = {"output"};
/* 输入 tensor */
int64_t input_shape[] = {1, 3, 640, 640};
int input_numel = 1 * 3 * 640 * 640;
float *input_data_host = new float[input_numel];
auto input_tensor = Ort::Value::CreateTensor(mem, input_data_host, input_numel, input_shape, 4);
/* 向输入 tensor 中填充数据 */
/* 输出 tensor */
int output_numbox = output_dims[1];
int output_numprob = output_dims[2];
int output_numel = 1 * output_numbox * output_numprob;
float *output_data_host = new float[output_numel];
int64_t output_shape[] = {1, output_numbox, output_numprob};
auto output_tensor = Ort::Value::CreateTensor(mem, output_data_host, output_numel, output_shape, 3);
// 模型推理
Ort::RunOptions options;
session.Run(options,
(const char *const *)input_names, &input_tensor, 1,
(const char *const *)output_names, &output_tensor, 1);
/* python run方法
res = session.run(["output"]，{"images" : input_tensor})
参数一：list，元素是输出节点；
参数二：dict，key 是输入节点名称 value 是输入的 tensor；
返回值：list，每个输出节点的输出值；
*/
/* 结果解析 */
// 释放资源
delete[] input_data_host;
delete[] output_data_host;
}OpenVINO推理
本节简要描述利用 OpenVINO 部署 ONNX 模型过程，其具体流程为：
- 初始化 OpenVINO 核心：创建一个
ov::Core对象，这是模型加载和推理的基础。 - 加载和编译模型：使用
core.compile_model方法加载 ONNX 模型文件，并将其编译到默认设备。 - 创建推理请求：调用
model.create_infer_request创建推理请求实例。 - 获取输入输出张量：使用
get_input_tensor和get_output_tensor方法获取输入和输出张量。如果模型支持动态形状（dynamic shape），需要显式设置输入张量的形状。 - 准备输入数据：调用
input.data()获取输入张量的内存指针。使用常规方式（如读取图像、归一化等）填充输入数据。 - 模型推理。
- 输出结果解析：获取输出张量的数据指针。根据模型的输出格式解析结果。
- 释放内存。
void onenvino_infer() {
ov::Core core;
// 编译模型
auto model = core.compile_model("../src/cuda-tensorrt-senior-api/static/yolov5s_modified.onnx");
// 创建推理请求
auto iq = model.create_infer_request();
// 获取输入与输出
auto input = iq.get_input_tensor(0);
auto output = iq.get_output_tensor(0);
// 动态 shape 需要手动确定输入维度
input.set_shape({1, 3, 640, 640});
// 注意：此处一定要先设置输入维度之后在获取输入 data 指针，否则获取到的 data 为空
float *input_data_host = input.data<float>();
/* 输入数据 */
iq.infer(); // 模型推理
float *output_data_host = output.data<float>(); // 获取输出 data 指针
/* 结果解析 */
std::cout << "Ending!" << std::endl;
}
