Keras到TensorRT部署
简述
-
TensorRT是什么 TensorRT是一个高性能的深度学习推理(Inference,相当于predict)优化器,可以为深度学习应用提供低延迟、高吞吐率的部署推理。可以认为TensorRT是一个只有前向传播的深度学习框架,这个框架可以将TensorFlow的网络模型解析,然后与TensorRT中对应的层进行一一映射、统一转换到TensorRT中,针对NVIDIA自家GPU实施优化策略,并进行部署加速。
-
学习TensorRT遇到的困难 初步接触TensorRT时发现,在各大博客上对此框架仅有简略介绍、没有较详细的开发教程资料。此框架需要使用C++语言编写,在环境搭建和代码测试上我们毫无经验,需要不断试错。对于代码报错,网上极少有对应的可行解决方案借鉴。最后我们学习英伟达官方提供的样例,模仿搭建了一个能够实现解析模型结构和权重、调用GPU的CUDA进行预测等功能的C++工程。
-
开发进度和流程优化 经过不断试错,我们成功开发了符合我们需求的TensorRT C++工程,并且融入了OpenCV图片处理、自定义HWC-CHW通道转换函数、语义分割结果的可视化处理等模块功能。同时发现,可以通过将keras框架的h5模型文件转换为ONNX格式喂入RT,但是存在网络结构解析较慢的情况;网上有没有可以加速模型解析的方案,但是通过研究英伟达官方提供的案例,我们尝试了将onnx转换为engine引擎的隐藏方法,省去了模型解析的步骤,使得模型可以直接加载到显存之中,极大压缩了程序加载时间。
格式转换
注意点:
to_channel_first指的是是否需要把NHWC通道转换为NCHW
TensorRT 5.1.5.0 GA 对于 ONNX target_opset=7 兼容性较好
import warnings,os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
warnings.filterwarnings("ignore")
# in tfrt 2.1.0
# pip install keras2onnx
import tensorflow as tf
import keras2onnx
import onnx
from tensorflow.keras.optimizers import Adam
'''
注意点:
to_channel_first指的是是否需要把NHWC通道转换为NCHW
TensorRT 5.0 GA 对于 ONNX target_opset=7 兼容性较好
'''
def h52onnx(h5_model_path, onnx_model_path, to_channel_first):
opt = Adam(lr=1e-4)
def get_lr_metric(optimizer): # printing the value of the learning rate
def lr(y_true, y_pred):
return optimizer.lr
return lr
lr_metric = get_lr_metric(opt)
model = tf.keras.models.load_model(h5_model_path)#, {'lr': lr_metric})
# model.summary()
inputLayerName = model.get_layer(index=0).name
print('First Layer is', inputLayerName)
if(not to_channel_first):
inputLayerName = None
onnx_model = keras2onnx.convert_keras(model, '',target_opset=7,
channel_first_inputs=inputLayerName,
doc_string='Straka\'s model.')
onnx.save_model(onnx_model, onnx_model_path)
print('onnx saved.')
def onnx2engine(onnxPath, enginePath):
cmd = r'D:\TensorRT\TensorRT-5.1.5.0\bin\trtexec --verbose --onnx='
cmd += onnxPath
cmd += r' --saveEngine='
cmd += enginePath
# print(cmd)
# os.system(cmd)
# print('engine saved.\n')
result = os.popen(cmd)
context = result.read()
result.close()
if(context.splitlines()[-1][5:11]=='FAILED'):
print('生成引擎失败.\n')
print(context)
elif(context.splitlines()[-1][5:11]=='PASSED'):
print('生成引擎成功.\n')
return context
修改C++工程
以英伟达官方的 sampleOnnxMNIST.cpp为例,加入OpenCV和自定义序列化引擎的导入 。B站讲解视频。
- 图像输入处理(包括HWC->CHW通道转换和BGR->RGB转换还有BGR->GRAY)
// ==============Pre Process=============>
void idxTransformParall(std::vector<unsigned char>* in_file, std::vector<float>* out_file,
unsigned long start_h, unsigned long length, unsigned long image_h, unsigned long image_w,
unsigned long start, float (*pFun)(const unsigned char&), bool HWC) {
if (HWC) {
// HWC and BRG=>RGB
for (unsigned long h = start_h; h < start_h + length; ++h) {
for (unsigned long w = 0; w < image_w; ++w) {
(*out_file)[start + h * image_w * 3 + w * 3 + 0] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 2]);
(*out_file)[start + h * image_w * 3 + w * 3 + 1] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 1]);
(*out_file)[start + h * image_w * 3 + w * 3 + 2] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 0]);
}
}
}
else {
// CHW and BRG=>RGB
for (unsigned long h = start_h; h < start_h + length; ++h) {
for (unsigned long w = 0; w < image_w; ++w) {
(*out_file)[start + 0 * image_h * image_w + h * image_w + w] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 2]);
(*out_file)[start + 1 * image_h * image_w + h * image_w + w] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 1]);
(*out_file)[start + 2 * image_h * image_w + h * image_w + w] =
(*pFun)((*in_file)[h * image_w * 3 + w * 3 + 0]);
}
}
}
}
// 图片 高 宽 按比例缩放 数值处理函数 通道模式 线程数
std::vector<float> imagePreprocess(const std::vector<cv::Mat>& images, const int& image_h, const int& image_w, bool is_padding, float(*pFun)(const unsigned char&), bool HWC, int worker) {
// image_path ===> cv::Mat ===> resize(padding) ===> CHW/HWC (BRG=>RGB)
// 测试发现RGB转BGR的cv::cvtColor 和 HWC 转 CHW非常耗时,故将其合并为一次操作
const unsigned long image_length = image_h * image_w * 3;
std::vector<float> fileData(images.size() * image_length);
for (unsigned long img_count = 0; img_count < images.size(); ++img_count) {
cv::Mat image = images[img_count].clone();
cv::Mat prodessed_image(image_h, image_w, CV_8UC3);
if (is_padding) {
int ih = image.rows;
int iw = image.cols;
float scale = std::min(static_cast<float>(image_w) / static_cast<float>(iw), static_cast<float>(image_h) / static_cast<float>(ih));
int nh = static_cast<int>(scale * static_cast<float>(ih));
int nw = static_cast<int>(scale * static_cast<float>(iw));
int dh = (image_h - nh) / 2;
int dw = (image_w - nw) / 2;
cv::Mat resized_image(nh, nw, CV_8UC3);
cv::resize(image, resized_image, cv::Size(nw, nh));
cv::copyMakeBorder(resized_image, prodessed_image, dh, image_h - nh - dh, dw, image_w - nw - dw, cv::BORDER_CONSTANT, cv::Scalar(128, 128, 128));
}
else {
cv::Mat resized_image(image_h, image_w, CV_8UC3);
cv::resize(image, prodessed_image, cv::Size(image_w, image_h));
}
std::vector<unsigned char> file_data = prodessed_image.reshape(1, 1);
// 并发
unsigned long min_threads;
if (worker < 0) {
const unsigned long min_length = 64;
min_threads = (image_h - 1) / min_length + 1;
}
else if (worker == 0) {
min_threads = 1;
}
else {
min_threads = worker;
}
const unsigned long cpu_max_threads = std::thread::hardware_concurrency();
const unsigned long num_threads = std::min(cpu_max_threads != 0 ? cpu_max_threads : 1, min_threads);
const unsigned long block_size = image_h / num_threads;
std::vector<std::thread> threads(num_threads - 1);
unsigned long block_start = 0;
for (auto& t : threads) {
t = std::thread(idxTransformParall, &file_data, &fileData, block_start, block_size, image_h, image_w, img_count * image_length, pFun, HWC);
block_start += block_size;
}
idxTransformParall(&file_data, &fileData, block_start, image_h - block_start, image_h, image_w, img_count * image_length, pFun, HWC);
for (auto& t : threads) {
t.join();
}
}
return fileData;
}
//=========================main==========================
startTime = clock(); //程序开始计时
float data[INPUT_H * INPUT_W * INPUT_C];
// 以单通道灰度格式读入
cv::ImreadModes mode = cv::ImreadModes::IMREAD_GRAYSCALE;
if (INPUT_C == 3)
// 以3通道彩色格式读入
mode = cv::ImreadModes::IMREAD_COLOR;
cv::Mat img = cv::imread(locateFile(imgPath, gArgs.dataDirs), mode);
gLogInfo << imgPath << " imgRawSize = " << img.size << std::endl;
if (INPUT_C == 3){
float(*pFunc)(const unsigned char&);
pFunc = [](const unsigned char& x) {return static_cast<float>(x) / 255; };
auto* pFuncc = pFunc;
std::vector<float> imgData;
imgData = imagePreprocess(std::vector<cv::Mat>{img}, INPUT_H, INPUT_W, doPadding, pFuncc, isHWC, 16);
for (int i = 0; i < INPUT_H * INPUT_W * INPUT_C; i++)
data[i] = (float)imgData[i];
//cv::cvtColor(img, img, cv::COLOR_BGR2RGB); // 改变颜色通道
//cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), cv::INTER_CUBIC); //三次插值缩放
//for (int i = 0; i < INPUT_H * INPUT_W * INPUT_C; i++)
// data[i] = float(img.data[i]) / 255.0f;
}
else {
cv::resize(img, img, cv::Size(INPUT_W, INPUT_H), cv::INTER_CUBIC); //三次插值缩放
//if (INPUT_C == 3)
//cv::cvtColor(img, img, cv::COLOR_BGR2RGB); // 改变颜色通道
for (int i = 0; i < INPUT_H * INPUT_W * INPUT_C; i++)
data[i] = float(img.data[i]) / 255.0f;
}
endTime = clock(); //程序结束用时
gLogInfo << "img dealt, using " << (double)(endTime - startTime) / CLOCKS_PER_SEC << "s." << std::endl;
- 序列化引擎
作用:当加载ONNX模型较慢时,可以先把ONNX转格式为ENGINE,然后直接加载引擎进行推理。
ICudaEngine* loadEngine(const std::string& engine, int DLACore, std::ostream& err, IRuntime* runtime)
{
IBuilder* builder = createInferBuilder(gLogger.getTRTLogger());
assert(builder != nullptr);
nvinfer1::INetworkDefinition* network = builder->createNetwork();
auto parser = nvonnxparser::createParser(*network, gLogger.getTRTLogger());
std::ifstream engineFile(engine, std::ios::binary);
if (!engineFile)
{
err << "Error opening engine file: " << engine << std::endl;
return nullptr;
}
engineFile.seekg(0, engineFile.end);
long int fsize = engineFile.tellg();
engineFile.seekg(0, engineFile.beg);
std::vector<char> engineData(fsize);
engineFile.read(engineData.data(), fsize);
if (!engineFile)
{
err << "Error loading engine file: " << engine << std::endl;
return nullptr;
}
if (DLACore != -1)
{
runtime->setDLACore(DLACore);
}
ICudaEngine* Engine = runtime->deserializeCudaEngine(engineData.data(), fsize, nullptr);
Engine->serialize(); //!!!序列化引擎!!!
parser->destroy();
network->destroy();
builder->destroy();
return Engine;
}
bool saveEngine(const ICudaEngine& engine, const std::string& fileName, std::ostream& err)
{
std::ofstream engineFile(fileName, std::ios::binary);
if (!engineFile)
{
err << "Cannot open engine file: " << fileName << std::endl;
return false;
}
IHostMemory* serializedEngine{ engine.serialize() };
if (serializedEngine == nullptr)
{
err << "Engine serialization failed" << std::endl;
return false;
}
engineFile.write(static_cast<char*>(serializedEngine->data()), serializedEngine->size());
return !engineFile.fail();
}
//in main
engine = loadEngine(locateFile(enginePath, gArgs.dataDirs), gArgs.useDLACore, gLogError, runtime);
- 结果输出
if (saveOutputImg) {
//输出语义分割图片
cv::Mat imageOut = cv::Mat::zeros(INPUT_H, INPUT_W, CV_8UC3);
for (int i = 0; i < INPUT_H; ++i) {
for (int j = 0; j < INPUT_W; ++j) {
imageOut.at<cv::Vec3b>(i, j)[0] = int(prob[i * INPUT_W + j] * 255);
imageOut.at<cv::Vec3b>(i, j)[1] = int(prob[i * INPUT_W + j] * 255);
imageOut.at<cv::Vec3b>(i, j)[2] = int(prob[i * INPUT_W + j] * 255);
}
}
cv::imwrite((locateFile(imgPath, gArgs.dataDirs).substr(0, (locateFile(imgPath, gArgs.dataDirs).length() - imgPath.length())) + "result.jpg"), imageOut);
gLogInfo << "result saved." << std::endl;
}
- 有报错stack flow
vs的项目属性->链接器->系统->堆栈保留大小,设置到160000000,可支持2000*2000*3
- 查看模型信息
//mInputDims = Engine->getBindingDimensions(Engine->getBindingIndex(inputLayerName));
//mOutputDims = Engine->getBindingDimensions(Engine->getBindingIndex(outputLayerName));
//gLogInfo << "getBindingIndex " << Engine->getBindingIndex(inputLayerName) << " " << Engine->getBindingIndex(outputLayerName) << std::endl;
//gLogInfo << "getNbLayers " << Engine->getNbLayers() << std::endl;
//gLogInfo << "getNbBindings " << Engine->getNbBindings() << std::endl;
//for (int i = 0; i < Engine->getNbBindings(); i++) {
// if(Engine ->bindingIsInput(i))
// mInputDims = Engine->getBindingDimensions(i);
// else
// mOutputDims = Engine->getBindingDimensions(i);
//}
//gLogInfo << "Input Dims = " << mInputDims.nbDims << std::endl;
//for (int i = 0; i < mInputDims.nbDims; i++) {
// gLogInfo << "Dim " << i << " includes " << mInputDims.d[i] << std::endl;
//}
//gLogInfo << "Output Dims = " << mOutputDims.nbDims << std::endl;
//for (int i = 0; i < mOutputDims.nbDims; i++) {
// gLogInfo << "Dim " << i << " includes " << mOutputDims.d[i] << std::endl;
//}