使用Tensorrt部署，C++ API yolov7_pose模型
虽然标题叫部署yolov7_pose模型，但是接下来的教程可以使用Tensorrt部署任何pytorch模型。
仓库地址：https://github.com/WongKinYiu/yolov7/tree/pose
系统版本：ubuntu18.4
驱动版本：CUDA Version: 11.4
在推理过程中，基于 TensorRT 的应用程序的执行速度可比 CPU 平台的速度快 40 倍。借助 TensorRT，您可以优化在所有主要框架中训练的神经网络模型，精确校正低精度，并最终将模型部署到超大规模数据中心、嵌入式或汽车产品平台中。
TensorRT 以 NVIDIA 的并行编程模型 CUDA 为基础构建而成，可帮助您利用 CUDA-X 中的库、开发工具和技术，针对人工智能、自主机器、高性能计算和图形优化所有深度学习框架中的推理。
TensorRT 针对多种深度学习推理应用的生产部署提供 INT8 和 FP16 优化，例如视频流式传输、语音识别、推荐和自然语言处理。推理精度降低后可显著减少应用延迟，这恰巧满足了许多实时服务、自动和嵌入式应用的要求。
我们部署的主要步骤为：将PytorchModel转化为OnnxModel，在将OnnxModel转化为TensorrtModel.
虽然看似步骤简单，但是坑还是有点多。
1.安装TensorRT
首先查看自己的Cuda版本，Windows 在cmd中执行nvidia-smi,Ubuntu在终端执行nvidia-smi即可查看cuda的版本。一般我们选择自己所能下载的最新的版本，避免有的算子没有实现的问题。我之前在这里被坑了一天。
然后根据版本在官网下载，点击Download，没有注册英伟达账号的需要注册账号登陆。官网地址：https://developer.nvidia.com/tensorrt
同意协议，然后根据自己的cuda版本选择，合适的版本。比如我的版本是cuda 11.4,一般选择Tar包
接下来将tar包或者zip包解压到你想安装的位置。这个软件解压即用，不用再安装。我们需要做的就是把软件的bin目录添加到环境变量。
Ubuntu:用vim打开～/.bashrc,将下面两行添加到文件最后面。
export LD_LIBRARY_PATH=/home/ubuntu/mySoftware/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
export PATH=/home/ubuntu/mySoftware/TensorRT-8.6.1.6/bin:$PATH
其中tensorrt的地址应该换成你解压的地址。然后sourse一下当前的终端
source ~/.bashrc
然后直接执行trtexec,如果没有报错证明成功安装了tensorrt
~/Downloads trtexec
&&&& RUNNING TensorRT.trtexec [TensorRT v8601] # trtexec -h
=== Model Options ===
--uff=<file> UFF model
--onnx=<file> ONNX model
--model=<file> Caffe model (default = no model, random weights used)
--deploy=<file> Caffe prototxt file
--output=<name>[,<name>]* Output names (it can be specified multiple times); at least one output is
......
2.转换pytorch模型为onnx格式的模型
先说yolo项目：项目目录下有个model/export.py
打开文件查看参数可以看到有一下参数设置。
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes') # ONNX-only
parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
parser.add_argument('--export-nms', action='store_true', help='export the nms part in ONNX model') # ONNX-only, #opt.grid has to be set True for nms export to work
opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
print(opt)
set_logging()
t = time.time()
根据自己模型设置合适的参数，注意如果你修改过模型的输出分类数，关键点数目。那么在导出nms层的时候就需要你自己手动修改网络模型。在models/common.py中的第361行修改non_max_suppression参数
class NMS(nn.Module):
# Non-Maximum Suppression (NMS) module
iou = 0.45 # IoU threshold
classes = None # (optional list) filter by class
def __init__(self, conf=0.25, kpt_label=False):
super(NMS, self).__init__()
self.conf=conf
self.kpt_label = kpt_label
def forward(self, x):
return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes, kpt_label=self.kpt_label,nc=2,nkpt=3)
class NMS_Export(nn.Module):
# Non-Maximum Suppression (NMS) module used while exporting ONNX model
iou = 0.45 # IoU threshold
classes = None # (optional list) filter by class
def __init__(self, conf=0.001, kpt_label=False):
super(NMS_Export, self).__init__()
self.conf = conf
self.kpt_label = kpt_label
def forward(self, x):
return non_max_suppression_export(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes, kpt_label=self.kpt_label,nc=2)
我们需要把nc和nkpt改为自己的设置的参数，比如我的分类为2，关键点数量为3。然后导出模型。
python --img-size 960 --weights /home/ubuntu/GITHUG/yolov7_pose/runs/train/exp2/weights/best.pt --grid --export-nms --simplify
如果顺利的话，我们会得到一个onnx格式的模型。我们可以打开https://netron.app/ 然后选择onnx模型打开。我们可以看到模型的图像
我们需要关注的就是模型的输入，输出。以及他们的形状。
从图中可以看出我的模型输入为images，大小为13 * 960 960输出为detections形状暂时不清楚。如果不清楚我们可以用onnxruntime跑一下查看形状
import onnxruntime
import numpy as np
import cv2
# 指定你的 ONNX 模型文件路径
onnx_model_path = '/home/ubuntu/GITHUG/yolov7_pose/runs/train/exp2/weights/best.onnx'
# 创建 ONNX Runtime 的推理会话
sess = onnxruntime.InferenceSession(onnx_model_path)
# 获取输入名称和形状
input_name = sess.get_inputs()[0].name
input_shape = sess.get_inputs()[0].shape
# 指定图像文件路径
image_path = '/home/ubuntu/GITHUG/yolov7_pose/501_png.rf.9cc0a917ca7972be6c8088aa9d17d651.jpg'
# 使用 OpenCV 读取图像
image = cv2.imread(image_path)
# 将图像调整为模型的输入形状
resized_image = cv2.resize(image, (input_shape[3], input_shape[2]))
# 将图像转换为浮点数并进行归一化
input_data = resized_image.astype(np.float32) / 255.0
# 将图像数据转换为 ONNX 模型期望的输入形状
input_data = np.transpose(input_data, [2, 0, 1])
input_data = np.expand_dims(input_data, axis=0)
# 运行推理
outputs = sess.run(None, {input_name: input_data})
# 输出模型的每个输出
for i, output_data in enumerate(outputs):
print(f"Output {i + 1}: {output_data}")
print(f"Output {output_data.shape}")
Output 1: [[8.01661621e+02 1.53809937e+02 9.72689453e+02 3.77949707e+02
4.21597920e-02 0.00000000e+00 5.15294671e-02 8.84101624e+02
2.51692810e+02 9.91612077e-01 9.03469177e+02 1.68072296e+02
6.35425091e-01 8.85691345e+02 1.72709320e+02 7.30822206e-01]
[7.85901917e+02 1.61294067e+02 9.64655701e+02 3.66809448e+02
4.08335961e-02 1.00000000e+00 6.32926583e-01 8.77714966e+02
2.57085205e+02 9.89280879e-01 8.91954224e+02 1.80863663e+02
2.32283741e-01 8.78342041e+02 1.87161697e+02 5.20734370e-01]
[7.05231201e+02 3.90309601e+02 7.51886230e+02 4.35935760e+02
1.86153594e-02 1.00000000e+00 6.94520175e-01 7.35046814e+02
4.11621490e+02 7.23196447e-01 7.14923584e+02 4.14582092e+02
4.62090850e-01 7.09832214e+02 4.12042603e+02 2.80124098e-01]
[4.01937828e+01 4.64705994e+02 1.51267151e+02 6.35167419e+02
1.55489137e-02 1.00000000e+00 9.99976933e-01 8.51227875e+01
5.72096252e+02 9.97074127e-01 8.59449158e+01 4.89000427e+02
9.83235717e-01 8.48072968e+01 5.18143494e+02 9.95639443e-01]
[4.67657043e+02 2.47014786e+02 6.09315125e+02 4.11179565e+02
1.50994565e-02 0.00000000e+00 1.29642010e-01 5.45577820e+02
3.71885773e+02 9.93896723e-01 5.56157104e+02 3.50972717e+02
9.97142434e-01 5.54454590e+02 3.20836670e+02 9.76849675e-01]
[3.69356445e+02 1.81159134e+01 4.91651611e+02 1.81579437e+02
1.44530777e-02 1.00000000e+00 9.98439074e-01 4.16761169e+02
1.16163483e+02 9.97292042e-01 4.29588745e+02 2.69206352e+01
9.79286790e-01 4.28487366e+02 8.01969910e+01 9.97563720e-01]
[7.12836548e+02 3.89805634e+02 7.66137817e+02 4.36001556e+02
1.32421134e-02 0.00000000e+00 2.13130921e-01 7.40284363e+02
4.09640594e+02 7.56286979e-01 7.18195129e+02 4.12563293e+02
1.05279446e-01 7.11785156e+02 4.14483521e+02 1.00254148e-01]
[7.01546204e+02 3.92902222e+02 7.31227966e+02 4.25415100e+02
1.30005283e-02 1.00000000e+00 9.94012475e-01 7.22401733e+02
4.12053406e+02 4.85429347e-01 7.12319214e+02 4.13364197e+02
7.06610680e-01 7.13084656e+02 4.11362488e+02 4.67233360e-01]
[6.80663696e+02 4.66796997e+02 7.09215454e+02 4.98112915e+02
1.06324852e-02 0.00000000e+00 6.49383068e-02 6.97597473e+02
4.87214142e+02 9.42029715e-01 6.90804749e+02 4.85028137e+02
9.82081532e-01 6.85866089e+02 4.70633820e+02 9.92424369e-01]]
Output (9, 16)
最后输出可以看出我的输出为1* 9 * 16，因为经过nms层后最后检测框的数量是不固定的所以应该是1 * x *16。仔细观察16纬的数据可以发现，每个数据都是
[x1,y1,x2,y2,confi,prob1,prob2,kpt1,conf1,pkt2,conf2,kpt3,conf3]
其中前四个数据为检测框，然后是置信度，分类概率，关键点以及关键点的置信度。
3.将onnx格式的模型转为.engine的tensorrt模型。
直接执行命令，然后等待模型转换成功。
trtexec --onnx=yolov7.onnx --fp16 --saveEngine=yolov7.engine
如果报错，比如什么算子不支持可以尝试更新tensorrt到最新版本。
4.C++部署
#include <iostream>
#include <fstream>
#include <vector>
#include <opencv2/opencv.hpp>
#include <NvInfer.h>
#include <cuda_runtime_api.h>
#define INPUT_W 960
#define INPUT_H 960
#define DEVICE 0 // GPU id
#define CONF_THRESH 0.2
using namespace nvinfer1;
class Logger : public ILogger {
void log(Severity severity, const char *msg) noexcept override {
// suppress info-level messages
if (severity <= Severity::kWARNING)
std::cout << msg << std::endl;
}
} logger;
#define CHECK(status) \
do\
{\
auto ret = (status);\
if (ret != 0)\
{\
std::cerr << "Cuda failure: " << ret << std::endl;\
abort();\
}\
} while (0)
float *blobFromImage(cv::Mat &img) {
cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
float *blob = new float[img.total() * 3];
int channels = 3;
int img_h = img.rows;
int img_w = img.cols;
for (int c = 0; c < channels; c++) {
for (int h = 0; h < img_h; h++) {
for (int w = 0; w < img_w; w++) {
blob[c * img_w * img_h + h * img_w + w] =
(((float) img.at<cv::Vec3b>(h, w)[c]) / 255.0f);
}
}
}
return blob;
}
cv::Mat static_resize(cv::Mat &img) {
float r = std::min(INPUT_W / (img.cols * 1.0), INPUT_H / (img.rows * 1.0));
int unpad_w = r * img.cols;
int unpad_h = r * img.rows;
cv::Mat re(unpad_h, unpad_w, CV_8UC3);
cv::resize(img, re, re.size());
cv::Mat out(INPUT_W, INPUT_H, CV_8UC3, cv::Scalar(114, 114, 114));
re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows)));
return out;
}
const char *INPUT_BLOB_NAME = "images";
const char *OUTPUT_BLOB_NAME = "detections";
static Logger gLogger;
static constexpr int MAX_OUTPUT_BBOX_COUNT = 100;
static constexpr int CLASS_NUM = 2;
static constexpr int LOCATIONS = 4;
static constexpr int KEY_POINTS_NUM = 3;
struct Keypoint {
float x;
float y;
float kpt_conf;
};
struct alignas(float) Detection {
//center_x center_y w h
float bbox[LOCATIONS];
float conf; // bbox_conf * cls_conf
float prob[CLASS_NUM]; // Probabilities for each class
// 3 keypoints
Keypoint kpts[KEY_POINTS_NUM];
};
void
doInference(IExecutionContext &context, float *input, float *output, const int output_size, const int input_shape) {
const ICudaEngine &engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void *buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
assert(engine.getBindingDataType(inputIndex) == nvinfer1::DataType::kFLOAT);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
assert(engine.getBindingDataType(outputIndex) == nvinfer1::DataType::kFLOAT);
// int mBatchSize = engine.getMaxBatchSize();
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], input_shape * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], output_size * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, input_shape * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueueV2(buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], output_size * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
static constexpr int DETECTION_SIZE = sizeof(Detection) / sizeof(float);
static void
postprocess_decode(float *feat_blob, float prob_threshold,std::vector<Detection> &objects_map) {
for (int i = 0; i < MAX_OUTPUT_BBOX_COUNT; i++) {
int base_index = i * DETECTION_SIZE; // Calculate the base index for the current detection
if (feat_blob[base_index + LOCATIONS] <= prob_threshold)
continue;
Detection det;
// Copy the detection information from feat_blob to the Detection structure
memcpy(&det, &feat_blob[base_index], DETECTION_SIZE * sizeof(float));
objects_map.push_back(det);
}
}
int main() {
char *trtModelStream{nullptr};
cudaSetDevice(DEVICE);
size_t size{0};
const char *engine_file_path = "/home/ubuntu/GITHUG/yolov7_pose/yolov7.engine";
std::ifstream file(engine_file_path, std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size);
file.close();
}
// create a model using the API directly and serialize it to a stream
IRuntime *runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size);
assert(engine != nullptr);
IExecutionContext *context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// auto out_dims = engine->getBindingDimensions(1);
int input_size = 1 * 3 * 960 * 960;
int output_size = MAX_OUTPUT_BBOX_COUNT * 16 * 1;
static float *prob = new float[output_size];
const char *input_image_path = "/home/ubuntu/GITHUG/yolov7_pose/501_png.rf.9cc0a917ca7972be6c8088aa9d17d651.jpg";
cv::Mat img = cv::imread(input_image_path);
cv::Mat pr_img = static_resize(img);
float *blob;
// cv::imshow("Image", pr_img);
blob = blobFromImage(pr_img);
cv::waitKey(200);
// 关闭窗口
cv::destroyAllWindows();
doInference(*context, blob, prob, output_size, input_size);
std::vector<Detection> objects_map;
for (int i = 0; i < prob[0] && i < MAX_OUTPUT_BBOX_COUNT; i++) {
std::cout << ": " << prob[i] << std::endl;
}
postprocess_decode(prob, CONF_THRESH, objects_map);
float r_w = INPUT_W / (img.cols * 1.0);
float r_h = INPUT_H / (img.rows * 1.0);
cv::cvtColor(pr_img, pr_img, cv::COLOR_RGB2BGR);
for (const auto &det: objects_map) {
// Access other information in the Detection structure as needed
// Example: Print bbox coordinates
std::cout << " Bbox: ";
for (int i = 0; i < LOCATIONS; i++) {
std::cout << det.bbox[i] << " ";
}
float r = 0.0;
if (img.rows <= img.cols) {
r = r_w;
} else {
r = r_h;
}
cv::Point pt1(det.bbox[0]/r, det.bbox[1]/r);
cv::Point pt2(det.bbox[2]/r, det.bbox[3]/r);
cv::rectangle(img, pt1, pt2, cv::Scalar(0, 255, 0), 2);
cv::Point point1(det.kpts[0].x / r, det.kpts[0].y / r);
cv::Point point2(det.kpts[1].x / r, det.kpts[1].y / r);
cv::Point point3(det.kpts[2].x / r, det.kpts[2].y / r);
// 画线段
cv::line(img, point1, point2, cv::Scalar(0, 0, 255), 2); // Scalar 参数表示颜色，这里是红色 (B, G, R)
cv::line(img, point2, point3, cv::Scalar(255, 0, 0), 2); // Scalar 参数表示颜色，这里是红色 (B, G, R)
cv::imshow("Rectangle", img);
cv::waitKey(0);
std::cout << std::endl;
}
}
这是我的demo以及最后的效果。
其中的关键代码为解析模型输出的部分，大家可以参考一下