1. TensorRT 简介
TensorRT 是为了模型高效部署推理的优化器。
多GPU训练 → 单GPU部署推理:一般的深度学习项目,训练时为了加快速度,会用GPU分布式训练。但在部署推理时,为了降低成本,往往使用单个GPU机器甚至嵌入式平台(比如 NVIDIA Jetson)进行部署。问题在于:
-
部署端要有与训练时相同的深度学习环境,比如Torch、Caffe、TensorFlow等。我们在部署的时候还需要搭建环境,这非常费时。
-
由于训练的网络模型参数量大,导致部署端机器的推理速度慢、延迟高,这在高实时性的场合下非常不理想。
为了提高部署推理时的速度,通常采取的方法有:
- 训练模型的优化。
-
轻量级神经网络,比如 SqueezeNet、MobileNet 等。其过程是基于经典模型提出的新的模型结构,然后用这些轻量模型重新训练再部署。
-
网络压缩算法,比如 剪枝、量化、蒸馏 等。其过程是在原有模型上经过算法得到参数量少但可以高效表达的模型。同样也是训练、微调后再部署。
- 部署端的优化。上述算法都是对模型本身的改造,而TensorRT则是对训练好的模型在部署推理过程中的优化算法。将训练好的文件直接丢进TensorRT中,不再依赖深度学习框架 (Torch、Caffe、TensorFlow 等)。
2. TensorRT 工作流程
2.1 TensorRT 部署大致流程
TensorRT部署的大致流程:1. 导出网络定义以及相关权重;2. 解析网络定义以及相关权重; 3. 根据显卡算子构建出最优执行计划; 4. 将执行计划序列化存储; 5. 反序列化执行计划; 6. 进行推理。
步骤3可知,TensorRT与硬件绑定,因此部署过程中,硬件(显卡)和软件(驱动、cuda、cudnn)如果发生了变化,流程则需要重新走。
2.2 各种深度学习框架使用TensorRT的方式
-
Caffe、TensorFlow、MatLab训练出的模型。TensorRT已经将他们集成到了框架中,因此可以直接解析他们的网络模型。
-
Caffe2、Pytorch、MXNet、CNTK、Chainer训练出的模型。先要将模型转为ONNX的通用深度学习模型,然后对ONNX模型做解析。
2.3 TensorRT 优化方式
- 层间融合或张量融合:在部署端模型推理时,每一层的操作都是GPU完成的,但实际上GPU通过启动不同的CUDA核心来完成计算。CUDA核心计算张量时,大量的时间会花费在CUDA核心的启动和对每一层输入/输出张量的读写操作上。因此TensorRT对层间张量横向或纵向合并,使层的数量减少,读写速度也就更快。
- 例如:该图中,纵向层融合 将卷积(Conv)、偏置(Bais)和激活层(ReLU)合并为一个CBR层,横向层融合将结构相同的层合并为一个层。
- 数据精度校准:
-
精度调整:在训练神经网络时,张量都是32位浮点精度,一旦网络训练完成,在部署推理时不需要反向传播,完全可以降低数据精度,如FP16(半精度浮点数)或INT8(8位数整型)。更低的数据精度会使内存占用和延迟更低,模型体积更小。
-
自动化校准:由于数据精度下降导致信息丢失,性能下降,因此TensorRT会提供自动化校准,以最好的匹配性能将FP32精度的数据降低为FP16/INT8,最小化性能损失
-
内核自动调整:TensorRT可以针对不同的算法,不同的网络模型,不同的GPU平台,进行CUDA核的调整,以保证当前模型在特定平台上以最优性能计算。
-
动态张量显存:TensorRT指定显存,避免显存重复申请,减少内存占用,提高重复使用效率。
-
多数据流操作:TensorRT并行处理多个输入流的可扩展设计。
TensorRT的优化方式对应的构建期和运行期:
2.4 TensorRT 部署的具体步骤
3. Anaconda 环境下 TensorRT 部署
3.1 CUDA、CUDNN 安装
3.2 TensorRT 安装
step 1. 下载 TensorRT
找到与CUDA对应版本的TensorRT版本下载,并解压放在磁盘上某个位置:
step 2. 配置TensorRT系统环境变量
step 3. 安装TensorRT依赖
想要在某个环境下部署,那就到该环境下
- 激活该环境:
python
conda activate yolo
- 下载TensorRT的依赖:
到环境中安装:
python
pip install D:\TensorRT-8.6.1.6\python\tensorrt-8.6.1-cp38-none-win_amd64.whl
- 下载Pycuda:
到pycuda网站,加载下来对应cuda、python版本的pycuda版本:我们cuda为1.20,因此只能找到最高版本cuda116,同时python为3.8版本。
到环境中安装:
python
pip install D:\pycuda\pycuda-2022.1+cuda116-cp38-cp38-win_amd64.whl
step 4. 测试TensorRT样例
在环境中,进入到样例目录,并运行 sample.py
sample能给出预测结果,则代表运行成功:
3.3 TensorRT 部署
案例:Anaconda环境下,Pytorch框架训练YOLOv8.pt → ONNX → TensorRT 部署
找到了一篇YOLOv8部署超级好的文章!
step 1. 将pt文件转为onnx文件
Python
# Export the onnx/engine modeld
model=YOLO(r'D:\work\project\ultralytics-main\runs\segment\train3\weights\best.pt')
model.export(format='onnx')
step 2. 将onnx文件转为engine文件
法1:使用TensorRT的trtexec执行程序
打开cmd,进入到trtexec.exe所在目录:
python
cd D:\TensorRT-8.6.1.6\bin
输入以下命令,将onnx文件转为engine文件:
python
trtexec --onnx=xxx.onnx --saveEngine=xxx.engine --fp16
# --onnx onnx存放路径
# --saveEngine engine存放路径
# --fp16 开启float16精度的推理
# --int8 开启 int8精度的推理
法2:YOLOv8中有直接将pt文件转为engine文件的方式
Python
# Export the onnx/engine modeld
model=YOLO(r'D:\work\project\ultralytics-main\runs\segment\train3\weights\best.pt')
model.export(format='engine')
step 3. TensorRT部署推理
Python API 推理:
python
# 导入必用依赖
import tensorrt as trt
import pycuda.autoinit #负责数据初始化,内存管理,销毁等
import pycuda.driver as cuda #GPU CPU之间的数据传输
# 创建logger:日志记录器
logger = trt.Logger(trt.Logger.WARNING)
# 创建runtime并反序列化生成engine
with open("sample.engine", "rb") as f, trt.Runtime(logger) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
# 分配CPU锁页内存和GPU显存
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
# 创建cuda流
stream = cuda.Stream()
# 创建context并进行推理
with engine.create_execution_context() as context:
# 将输入数据从锁页内存拷贝到显存
cuda.memcpy_htod_async(d_input, h_input, stream)
# 进行推理
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
# 将输出结果从显存拷贝到锁页内存
cuda.memcpy_dtoh_async(h_output, d_output, stream)
# 同步流
stream.synchronize()
# 返回主内存的输出结果. 该数据等同于原始模型的输出数据
return h_output
C++ API 推理:
python
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "NvInferPlugin.h"
#include "logging.h"
#include <opencv2/opencv.hpp>
#include "utils.h"
#include <string>
using namespace nvinfer1;
using namespace cv;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 640;
static const int INPUT_W = 640;
static const int _segWidth = 160;
static const int _segHeight = 160;
static const int _segChannels = 32;
static const int CLASSES = 80;
static const int Num_box = 8400;
static const int OUTPUT_SIZE = Num_box * (CLASSES+4 + _segChannels);//output0
static const int OUTPUT_SIZE1 = _segChannels * _segWidth * _segHeight ;//output1
static const float CONF_THRESHOLD = 0.1;
static const float NMS_THRESHOLD = 0.5;
static const float MASK_THRESHOLD = 0.5;
const char* INPUT_BLOB_NAME = "images";
const char* OUTPUT_BLOB_NAME = "output0";//detect
const char* OUTPUT_BLOB_NAME1 = "output1";//mask
struct OutputSeg {
int id; //结果类别id
float confidence; //结果置信度
cv::Rect box; //矩形框
cv::Mat boxMask; //矩形框内mask,节省内存空间和加快速度
};
void DrawPred(Mat& img,std:: vector<OutputSeg> result) {
//生成随机颜色
std::vector<Scalar> color;
srand(time(0));
for (int i = 0; i < CLASSES; i++) {
int b = rand() % 256;
int g = rand() % 256;
int r = rand() % 256;
color.push_back(Scalar(b, g, r));
}
Mat mask = img.clone();
for (int i = 0; i < result.size(); i++) {
int left, top;
left = result[i].box.x;
top = result[i].box.y;
int color_num = i;
rectangle(img, result[i].box, color[result[i].id], 2, 8);
mask(result[i].box).setTo(color[result[i].id], result[i].boxMask);
char label[100];
sprintf(label, "%d:%.2f", result[i].id, result[i].confidence);
//std::string label = std::to_string(result[i].id) + ":" + std::to_string(result[i].confidence);
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = max(top, labelSize.height);
putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
}
addWeighted(img, 0.5, mask, 0.8, 1, img); //将mask加在原图上面
}
static Logger gLogger;
void doInference(IExecutionContext& context, float* input, float* output, float* output1, int batchSize)
{
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() == 3);
void* buffers[3];
// 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);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
const int outputIndex1 = engine.getBindingIndex(OUTPUT_BLOB_NAME1);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));//
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float)));
// cudaMalloc分配内存 cudaFree释放内存 cudaMemcpy或 cudaMemcpyAsync 在主机和设备之间传输数据
// cudaMemcpy cudaMemcpyAsync 显式地阻塞传输 显式地非阻塞传输
// 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, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
CHECK(cudaMemcpyAsync(output1, buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
CHECK(cudaFree(buffers[outputIndex1]));
}
int main(int argc, char** argv)
{
if (argc < 2) {
argv[1] = "../models/yolov8n-seg.engine";
argv[2] = "../images/bus.jpg";
}
// create a model using the API directly and serialize it to a stream
char* trtModelStream{ nullptr }; //char* trtModelStream==nullptr; 开辟空指针后 要和new配合使用,比如89行 trtModelStream = new char[size]
size_t size{ 0 };//与int固定四个字节不同有所不同,size_t的取值range是目标平台下最大可能的数组尺寸,一些平台下size_t的范围小于int的正数范围,又或者大于unsigned int. 使用Int既有可能浪费,又有可能范围不够大。
std::ifstream file(argv[1], std::ios::binary);
if (file.good()) {
std::cout << "load engine success" << std::endl;
file.seekg(0, file.end);//指向文件的最后地址
size = file.tellg();//把文件长度告诉给size
//std::cout << "\nfile:" << argv[1] << " size is";
//std::cout << size << "";
file.seekg(0, file.beg);//指回文件的开始地址
trtModelStream = new char[size];//开辟一个char 长度是文件的长度
assert(trtModelStream);//
file.read(trtModelStream, size);//将文件内容传给trtModelStream
file.close();//关闭
}
else {
std::cout << "load engine failed" << std::endl;
return 1;
}
Mat src = imread(argv[2], 1);
if (src.empty()) { std::cout << "image load faild" << std::endl; return 1; }
int img_width = src.cols;
int img_height = src.rows;
std::cout << "宽高:" << img_width << " " << img_height << std::endl;
// Subtract mean from image
static float data[3 * INPUT_H * INPUT_W];
Mat pr_img0, pr_img;
std::vector<int> padsize;
pr_img = preprocess_img(src, INPUT_H, INPUT_W, padsize); // Resize
int newh = padsize[0], neww = padsize[1], padh = padsize[2], padw = padsize[3];
float ratio_h = (float)src.rows / newh;
float ratio_w = (float)src.cols / neww;
int i = 0;// [1,3,INPUT_H,INPUT_W]
//std::cout << "pr_img.step" << pr_img.step << std::endl;
for (int row = 0; row < INPUT_H; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;//pr_img.step=widthx3 就是每一行有width个3通道的值
for (int col = 0; col < INPUT_W; ++col)
{
data[i] = (float)uc_pixel[2] / 255.0;
data[i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
data[i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.;
uc_pixel += 3;
++i;
}
}
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
bool didInitPlugins = initLibNvInferPlugins(nullptr, "");
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Run inference
static float prob[OUTPUT_SIZE];
static float prob1[OUTPUT_SIZE1];
//for (int i = 0; i < 10; i++) {//计算10次的推理速度
// auto start = std::chrono::system_clock::now();
// doInference(*context, data, prob, prob1, 1);
// auto end = std::chrono::system_clock::now();
// std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
// }
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, prob1, 1);
auto end = std::chrono::system_clock::now();
std::cout << "推理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
std::vector<int> classIds;//结果id数组
std::vector<float> confidences;//结果每个id对应置信度数组
std::vector<cv::Rect> boxes;//每个id矩形框
std::vector<cv::Mat> picked_proposals; //后续计算mask
// 处理box
int net_length = CLASSES + 4 + _segChannels;
cv::Mat out1 = cv::Mat(net_length, Num_box, CV_32F, prob);
start = std::chrono::system_clock::now();
for (int i = 0; i < Num_box; i++) {
//输出是1*net_length*Num_box;所以每个box的属性是每隔Num_box取一个值,共net_length个值
cv::Mat scores = out1(Rect(i, 4, 1, CLASSES)).clone();
Point classIdPoint;
double max_class_socre;
minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
max_class_socre = (float)max_class_socre;
if (max_class_socre >= CONF_THRESHOLD) {
cv::Mat temp_proto = out1(Rect(i, 4 + CLASSES, 1, _segChannels)).clone();
picked_proposals.push_back(temp_proto.t());
float x = (out1.at<float>(0, i) - padw) * ratio_w; //cx
float y = (out1.at<float>(1, i) - padh) * ratio_h; //cy
float w = out1.at<float>(2, i) * ratio_w; //w
float h = out1.at<float>(3, i) * ratio_h; //h
int left = MAX((x - 0.5 * w), 0);
int top = MAX((y - 0.5 * h), 0);
int width = (int)w;
int height = (int)h;
if (width <= 0 || height <= 0) { continue; }
classIds.push_back(classIdPoint.y);
confidences.push_back(max_class_socre);
boxes.push_back(Rect(left, top, width, height));
}
}
//执行非最大抑制以消除具有较低置信度的冗余重叠框(NMS)
std::vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, CONF_THRESHOLD, NMS_THRESHOLD, nms_result);
std::vector<cv::Mat> temp_mask_proposals;
std::vector<OutputSeg> output;
Rect holeImgRect(0, 0, src.cols, src.rows);
for (int i = 0; i < nms_result.size(); ++i) {
int idx = nms_result[i];
OutputSeg result;
result.id = classIds[idx];
result.confidence = confidences[idx];
result.box = boxes[idx]&holeImgRect;
output.push_back(result);
temp_mask_proposals.push_back(picked_proposals[idx]);
}
// 处理mask
Mat maskProposals;
for (int i = 0; i < temp_mask_proposals.size(); ++i)
maskProposals.push_back(temp_mask_proposals[i]);
Mat protos = Mat(_segChannels, _segWidth * _segHeight, CV_32F, prob1);
Mat matmulRes = (maskProposals * protos).t();//n*32 32*25600 A*B是以数学运算中矩阵相乘的方式实现的,要求A的列数等于B的行数时
Mat masks = matmulRes.reshape(output.size(), { _segWidth,_segHeight });//n*160*160
std::vector<Mat> maskChannels;
cv::split(masks, maskChannels);
Rect roi(int((float)padw / INPUT_W * _segWidth), int((float)padh / INPUT_H * _segHeight), int(_segWidth - padw / 2), int(_segHeight - padh / 2));
for (int i = 0; i < output.size(); ++i) {
Mat dest, mask;
cv::exp(-maskChannels[i], dest);//sigmoid
dest = 1.0 / (1.0 + dest);//160*160
dest = dest(roi);
resize(dest, mask, cv::Size(src.cols, src.rows), INTER_NEAREST);
//crop----截取box中的mask作为该box对应的mask
Rect temp_rect = output[i].box;
mask = mask(temp_rect) > MASK_THRESHOLD;
output[i].boxMask = mask;
}
end = std::chrono::system_clock::now();
std::cout << "后处理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
DrawPred(src, output);
cv::imshow("output.jpg", src);
char c = cv::waitKey(0);
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
system("pause");
return 0;
}