好的,阿杰,我们来看 TensorRT C++ 部署 的完整思路、步骤和注意事项。我会把流程拆解清楚,适合直接在工程中实践。
1️⃣ 前提条件
- 安装 NVIDIA 驱动(对应你的 GPU 型号和 CUDA 版本)
- 安装 CUDA Toolkit
- 安装 cuDNN
- 安装 TensorRT SDK(包含
libnvinfer、libnvinfer_plugin、include等) - C++ 编译器支持(Linux 推荐 GCC ≥ 7,Windows 推荐 MSVC 2019+)
参考官方安装指南:TensorRT Installation Guide
2️⃣ 基本部署流程
2.1 导出模型
- TensorRT 支持 ONNX、Caffe 等模型格式
- 对 PyTorch 模型:
import torch dummy_input = torch.randn(1, 3, 224, 224).cuda() torch.onnx.export(model, dummy_input, "model.onnx", opset_version=13) - 对 TensorFlow 模型:可先转换为
SavedModel再用tf2onnx转 ONNX
2.2 使用 TensorRT 生成 engine 文件
示例 C++ 代码:
#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <iostream>
#include <fstream>
using namespace nvinfer1;
int main() {
// 创建 Logger
class Logger : public ILogger {
void log(Severity severity, const char* msg) noexcept override {
if (severity <= Severity::kWARNING)
std::cout << msg << std::endl;
}
} logger;
// 创建 Builder
IBuilder* builder = createInferBuilder(logger);
INetworkDefinition* network = builder->createNetworkV2(0U);
auto parser = nvonnxparser::createParser(*network, logger);
if (!parser->parseFromFile("model.onnx", static_cast<int>(ILogger::Severity::kWARNING))) {
std::cerr << "ONNX parse failed!" << std::endl;
return -1;
}
// Builder 配置
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 30); // 1GB
config->setFlag(BuilderFlag::kFP16); // 开启 FP16 加速
// 生成 engine
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
if (!engine) {
std::cerr << "Build engine failed!" << std::endl;
return -1;
}
// 保存 engine
IHostMemory* modelStream = engine->serialize();
std::ofstream file("model.trt", std::ios::binary);
file.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
// 释放资源
modelStream->destroy();
engine->destroy();
network->destroy();
builder->destroy();
return 0;
}
参考官方文档:TensorRT C++ API
2.3 推理部署
加载 engine 并执行推理:
#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <vector>
int main() {
// 1. 反序列化 engine
std::ifstream file("model.trt", std::ios::binary);
file.seekg(0, std::ios::end);
size_t size = file.tellg();
file.seekg(0, std::ios::beg);
std::vector<char> engineData(size);
file.read(engineData.data(), size);
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), size);
// 2. 创建执行上下文
nvinfer1::IExecutionContext* context = engine->createExecutionContext();
// 3. 准备输入输出缓冲
void* buffers[2];
cudaMalloc(&buffers[0], inputSize);
cudaMalloc(&buffers[1], outputSize);
// 4. 将输入数据拷贝到 GPU
cudaMemcpy(buffers[0], hostInput, inputSize, cudaMemcpyHostToDevice);
// 5. 推理
context->executeV2(buffers);
// 6. 拷贝输出回 CPU
cudaMemcpy(hostOutput, buffers[1], outputSize, cudaMemcpyDeviceToHost);
// 7. 释放资源
cudaFree(buffers[0]);
cudaFree(buffers[1]);
context->destroy();
engine->destroy();
runtime->destroy();
return 0;
}
3️⃣ 部署优化建议
- 使用 FP16/INT8:减少显存和计算时间
- 批量推理:使用 batch,提高 GPU 利用率
- 使用 CUDA Stream:异步推理,提高吞吐
- TensorRT Plugin:自定义操作支持非标准层
- 序列化 engine:避免每次启动重新 build,提高启动速度
发表回复