人工智能深度学习计算机视觉【免费下载链接】tensorrtxImplementation of popular deep learning networks with TensorRT network definition API项目地址https://gitcode.com/gh_mirrors/te/tensorrtx点击查看免费下载导读本文以开源仓库 tensorrtx 中的 crnn 模块为主体完整讲解如何将 PyTorch 训练好的 CRNN卷积循环神经网络场景文字识别模型一步步转换为 TensorRT 引擎并部署推理从权重导出wts 文件、CMake 构建到纯 TensorRT Network Definition API 的 CNN 双向 LSTM 网络搭建再到 CTC 解码与结果验证。读完本文你将掌握一套无需解析器、完全基于 TensorRT 原生 API 手写网络的部署方法论可直接复用于其他 RNN/LSTM 类模型。一、CRNN 模型与 tensorrtx/crnn 的实现定位CRNNConvolutional Recurrent Neural Network是经典的场景文字识别STR方案结构上由 CNN 卷积主干、RNN循环层和 CTC 转录层三部分组成能够直接对整行不定长文本图像做序列识别。本项目tensorrtx即 Implementation of popular deep learning networks with TensorRT network definition API采用**网络定义 APINetwork Definition API**而非 ONNX/TensorFlow 解析器来重建 CRNN这意味着网络中的每一个卷积、池化、LSTM 单元都是通过nvinfer1接口逐层手动构建的。crnn 模块的核心文件构成文件作用crnn/crnn.cpp网络定义、引擎序列化/反序列化、推理与 CTC 解码主程序crnn/genwts.py将 PyTorch 权重crnn.pth导出为 TensorRT 可读的crnn.wtscrnn/CMakeLists.txt构建脚本链接 CUDA / OpenCV / TensorRT 库crnn/logging.hTensorRT 日志器封装仓库内多模块共用crnn/README.md官方使用说明本文所依据的主体文档从源码结构看模型的 PyTorch 参考实现为 meijieru/crnn.pytorch其权重和导出脚本即来源自该工程下文会详细说明依赖关系。二、整体运行流程三步从 PyTorch 权重到 TensorRT 推理依据 crnn/README.md 的 How to Run 说明完整流程分三步导出 wts → 构建引擎-s→ 反序列化推理-d。第 1 步从 PyTorch 生成 crnn.wts# 克隆本仓库与 crnn.pytorch此处给出 README 原始描述按需调整路径 git clone https://github.com/wang-xinyu/tensorrtx.git git clone https://github.com/meijieru/crnn.pytorch.git # 下载其预训练权重 crnn.pth # 将 tensorrtx/crnn/genwts.py 复制到 crnn.pytorch/ 目录下 # 进入 crnn.pytorch/ 目录执行 python genwts.py # 执行完成后会生成一个 crnn.wts 文件为什么 genwts.py 必须复制到 crnn.pytorch 目录下执行因为 crnn/genwts.py 开头直接import utils与import models.crnn as crnn这两个模块分别由 crnn.pytorch 仓库的utils.py与models/crnn.py提供包括 CRNN 类定义、LSTM 结构等脱离该目录无法运行。第 2 步构建 tensorrtx/crnn 并序列化引擎# 将 crnn.wts 放入 tensorrtx/crnn 目录 # 进入 tensorrtx/crnn 目录 mkdir build cd build cmake .. make sudo ./crnn -s # 序列化模型到 plan 文件即 crnn.engine # 将 crnn.pytorch/data/demo.png 复制到当前目录 sudo ./crnn -d # 反序列化 plan 文件并执行推理第 3 步检查输出README 给出的预期输出如下raw: a-----v--a-i-l-a-bb-l-e--- sim: available其中raw是未折叠的原始类别序列含 blank 符-sim是经过 CTC 折叠后的最终识别文本available两条输出含义详见本文第七节解码原理。三、构建环境与依赖解析CMakeLists.txtcrnn/CMakeLists.txt 明确了运行环境要求这也是让上面两条命令能跑通的前提C 标准add_definitions(-stdc11)与set(CMAKE_CXX_STANDARD 11)并以Debug类型构建CUDA 依赖find_package(CUDA REQUIRED)链接cudartCUDA 运行时TensorRT 依赖链接nvinferOpenCV 依赖find_package(OpenCV)并链接${OpenCV_LIBS}用于推理阶段的图片读取与预处理嵌入式平台适配当CMAKE_SYSTEM_PROCESSOR匹配aarch64时头文件与库路径切换为/usr/local/cuda/targets/aarch64-linux/...说明该实现同时兼容 Jetson 等 ARM 平台。因此部署环境需要预先安装好 CUDA、TensorRT 与 OpenCV 三套软件栈具体安装步骤可参考仓库中的 tutorials/install.md 与 tutorials/getting_started.md。构建时使用sudo通常是为了在目标 GPU 设备上正常申请显存与执行 CUDA 操作README 中的-s、-d两条命令均以sudo前缀给出请结合你的实际环境权限决定是否保留。四、第一步详解genwts.py 权重导出与 wts 文件格式4.1 导出脚本逻辑crnn/genwts.py 的核心逻辑非常简洁model_path ./data/crnn.pth model crnn.CRNN(32, 1, 37, 256) # 输入高度32、单通道、37类、隐藏层256 model.load_state_dict(torch.load(model_path)) image torch.ones(1, 1, 32, 100) # 构造一张 (1,1,32,100) 的占位输入 preds model(image) # 前向一次确保权重被正确加载随后脚本遍历model.state_dict()的所有键值对把每个张量 reshape 为一维后以大端序big-endian单精度浮点的十六进制形式逐元素写入文本文件crnn.wtsf.write({} {}\n.format(len(model.state_dict().keys()), ...)) for k, v in model.state_dict().items(): vr v.reshape(-1).cpu().numpy() f.write({} {}.format(k, len(vr))) for vv in vr: f.write( struct.pack(f, float(vv)).hex())注意crnn.CRNN(32, 1, 37, 256)中的四个参数分别对应输入图像高度 32、通道数 1、类别数 3736 个字符 1 个 blank以及 LSTM 隐藏单元数 256——这些数字必须与 C 端 crnn/crnn.cpp 中的网络定义严格一致否则加载权重时会出现张量尺寸不匹配。4.2 wts 文件格式crnn.wts是一个纯文本权重文件格式与 tensorrtx 其他模块统一可用 crnn/crnn.cpp 中的loadWeights函数印证首行一个整数表示权重张量的总数即state_dict的键数量其后每行[权重名] [元素个数] [以空格分隔的大端序 float 十六进制数据...]。C 侧读取时逐行解析按名存入std::mapstd::string, Weights供后续网络定义按cnn.conv0.weight、rnn.0.rnn.weight_ih_l0这类键名取用。如果文件缺失或路径错误loadWeights中的断言会直接提示Unable to load weight file. please check if the .wts file path is right!!!!!!五、第二步详解crnn.cpp 网络定义与引擎构建源码级5.1 网络输入输出与关键常量crnn/crnn.cpp 顶部定义了一组贯穿全流程的常量#define USE_FP16 // 注释掉此行则使用 FP32 #define DEVICE 0 // GPU id #define BATCH_SIZE 1 static const int INPUT_H 32; static const int INPUT_W 100; static const int OUTPUT_SIZE 26 * 37; // 26 个时间步 × 37 个类别 const char* INPUT_BLOB_NAME data; const char* OUTPUT_BLOB_NAME prob; const int ks[] {3, 3, 3, 3, 3, 3, 2}; // 每层卷积核尺寸 const int ps[] {1, 1, 1, 1, 1, 1, 0}; // 每层 padding const int ss[] {1, 1, 1, 1, 1, 1, 1}; // 每层 stride const int nm[] {64, 128, 256, 256, 512, 512, 512}; // 每层输出通道数 const std::string alphabet -0123456789abcdefghijklmnopqrstuvwxyz;这里的alphabet共 37 个字符索引 0 是 CTC 的 blank 占位符-随后是数字 0-9 和 26 个小写字母。OUTPUT_SIZE 26 * 37说明输入宽度 100 经过卷积/池化后在时间轴上保留 26 个序列位置每个位置输出 37 类的概率分布。5.2 CNN 主干7 层卷积 批归一化 池化createEnginecrnn/crnn.cpp首先注册输入张量ITensor* data network-addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});即单通道 32×100 的灰度图像数据布局为 CHW。随后按 CRNN 经典结构逐层堆叠关键调用的顺序为convRelu(..., 0)→ 最大池化2×2stride 2convRelu(..., 1)→ 最大池化2×2stride 2convRelu(..., 2, true)带 BatchNorm→convRelu(..., 3)→ 最大池化2×2stride (2,1)、padding (0,1)convRelu(..., 4, true)→convRelu(..., 5)→ 最大池化2×2stride (2,1)、padding (0,1)convRelu(..., 6, true)收尾。convRelu的实现crnn/crnn.cpp用addConvolutionNd建卷积层按ks/ss/ps/nm数组设置卷积核尺寸、步长、padding 和输出通道卷积后接addActivation(..., ActivationType::kRELU)做 ReLU 激活。第 2、4、6 层额外调用addBatchNorm2d挂接 BatchNorm。BatchNorm 折叠技巧TensorRT 的addBatchNorm2dcrnn/crnn.cpp并没有直接提供 BN 层而是利用公式将 BN 折叠进IScaleLayerper-channel 的 Scale 层scale[i] gamma[i] / sqrt(var[i] eps) shift[i] beta[i] - mean[i] * gamma[i] / sqrt(var[i] eps) power[i] 1.0其中gamma/beta来自.weight/.biasmean/var来自.running_mean/.running_vareps取1e-5。这一手法在 tensorrtx 的众多视觉模型中通用理解它有助于阅读仓库内其他模块如 resnet/resnet50.cpp的 BN 处理。5.3 时序维度变换与双向 LSTMCNN 输出经过 Shuffle 转置Permutation{1, 2, 0}后进入两层双向 LSTMcrnn/crnn.cppauto lstm0 addLSTM(network, weightMap, *sfl-getOutput(0), 256, rnn.0.rnn); // reshape 为 {26, 1, 1, 512} 后接 FC512 → 256embedding 层 // 再转置 reshape 回 {1, 26, 256}输入第二层 LSTM auto lstm1 addLSTM(network, weightMap, *sfl-getOutput(0), 256, rnn.1.rnn); // reshape 为 {26, 1, 1, 512} 后接 FC512 → 37作为最终输出addLSTMcrnn/crnn.cpp基于 TensorRT 的addRNNv2实现要点有二方向lstm-setDirection(RNNDirection::kBIDIRECTION)即双向 LSTM每层输出维度为隐藏单元数 × 2256 → 512权重拆分PyTorch 的 LSTM 权重是 4 个门input/forget/cell/output拼接在一起的整体张量因此splitLstmWeightscrnn/crnn.cpp会把weight_ih_l0等按count/4均分生成weight_ih_l00~weight_ih_l03再通过setWeightsForGate/setBiasForGate分别指定每个门的输入权重、循环权重与偏置正反向_l0与_l0_reverse各自独立设置。两层 LSTM 之后各接一个addFullyConnected全连接层第一层把 512 维压缩到 256对应 PyTorch 中的 embedding 层rnn.0.embedding.weight/bias第二层输出 37 维类别得分rnn.1.embedding.weight/bias。最终输出张量命名为prob并markOutput。5.4 引擎构建配置构建阶段crnn/crnn.cppbuilder-setMaxBatchSize(maxBatchSize); config-setMaxWorkspaceSize(16 * (1 20)); // workspace 上限 16MB #ifdef USE_FP16 config-setFlag(BuilderFlag::kFP16); #endif ICudaEngine* engine builder-buildEngineWithConfig(*network, *config);默认BATCH_SIZE 1通过-s序列化为crnn.engine二进制 plan 文件通过注释/取消注释USE_FP16宏即可在 FP16 与 FP32 精度间切换FP16 需要 GPU 支持可参考 tutorials/check_fp16_int8_support.md 检测。六、第三步详解反序列化、图像预处理与推理执行6.1 加载引擎-d分支crnn/crnn.cpp从crnn.engine读出字节流经createInferRuntime→deserializeCudaEngine→createExecutionContext完成引擎加载并断言绑定数为 2一个输入data、一个输出prob通过getBindingIndex获取输入/输出索引后分配 GPU 显存。6.2 输入预处理推理输入来自图片demo.pngREADME 要求从 crnn.pytorch/data 复制到 crnn 目录处理流程与 CRNN 训练约定一致crnn/crnn.cppcv::Mat img cv::imread(demo.png); cv::cvtColor(img, img, CV_BGR2GRAY); // 转灰度单通道 cv::resize(img, img, cv::Size(INPUT_W, INPUT_H)); // 强制缩放到 100×32 data[i] ((float)img.atuchar(i) / 255.0 - 0.5) * 2.0; // 归一化到 [-1, 1]如果demo.png缺失程序会打印demo.png not found !!!并退出。6.3 异步推理doInferencecrnn/crnn.cpp以 CUDA Stream 异步执行整条流水线cudaMemcpyAsync(buffers[0], input, ..., cudaMemcpyHostToDevice, stream); // 输入上载 context.enqueue(batchSize, buffers, stream, nullptr); // 异步推理 cudaMemcpyAsync(output, buffers[1], ..., cudaMemcpyDeviceToHost, stream); // 结果回传 cudaStreamSynchronize(stream);主程序还会用std::chrono统计一次推理的毫秒耗时并打印。七、输出解码raw 与 sim 的区别CTC 折叠原理推理得到的prob是 26×37 的概率矩阵。程序对每个时间步共 26 个做一次 argmax得到长度为 26 的类别索引序列preds再交给strDecodecrnn/crnn.cpp解码std::string strDecode(std::vectorint preds, bool raw) { if (raw) { for (auto v : preds) str.push_back(alphabet[v]); // raw直接逐字符映射 } else { for (size_t i 0; i preds.size(); i) { if (preds[i] 0 || (i 0 preds[i - 1] preds[i])) continue; str.push_back(alphabet[preds[i]]); // simCTC 折叠 } } }raw 模式不做任何处理直接把 26 个 argmax 索引映射为字母。由于 blank 符-在索引 0 处raw 输出会包含大量-例如a-----v--a-i-l-a-bb-l-e---sim 模式执行标准的CTC collapse两步操作——① 跳过索引 0blank 占位符② 合并相邻的连续重复字符。于是 raw 序列a-----v--a-i-l-a-bb-l-e---折叠为available。这正是 CTC 损失训练模型的标准解码行为也是校验部署正确性的最直接手段只要 sim 输出等于预期文本就说明权重加载、网络构建与推理链路完全正确。八、工程要点与常见问题权重路径createEngine中写死从../crnn.wts加载权重crnn/crnn.cpp因此按 README 流程应将crnn.wts放入tensorrtx/crnn目录并从build/子目录发起构建才能解析到上一级的 wts 文件图片路径推理阶段读取的是当前工作目录下的demo.pngcrnn/crnn.cpp务必在build/中先复制 demo.png 再执行./crnn -d精度开关USE_FP16宏控制 FP16/FP32FP16 可显著降低带宽与加速但需要硬件支持Tesla/RTX 等不支持时会退化为不可用需关闭该宏BatchSize 固定为 1BATCH_SIZE、输入输出缓冲区尺寸均在源码中硬编码如需批处理需同步修改main中static float data[...]等缓冲区大小及OUTPUT_SIZE相关计算OCR 文本为小写字母与数字模型字母表仅含-0123456789abcdefghijklmnopqrstuvwxyz输入文本需为小写超出字母表范围的字符不会被正确识别输入宽高约定网络输入固定为 32×100 灰度图预处理强制 resize宽度方向上的特征缩放会影响识别效果实际应用时应尽量保证输入文本比例接近训练分布。九、延伸阅读本文只聚焦 CRNN 一条完整链路但其中用到的wts 权重导出 → Network Definition API 建图 → 序列化/反序列化 → 后处理解码四段式方法在 tensorrtx 仓库中具有高度通用性同为序列/时序类模型可参考 tsm/README.md、ufld/README.md相同 wts 格式与权重加载方式可见于 lenet/README.md、resnet/README.md 等模块环境搭建与通用问题可查阅 tutorials/install.md、tutorials/faq.md、tutorials/getting_started.md。如果希望把 CRNN 接入真实 OCR 管线可将 sim 解码输出与词典/拼音纠错结合但请注意仓库本身只提供 TensorRT 推理实现与示例验证不包含完整的文本纠错与业务后处理这部分需要依据实际业务自行扩展。赞分享人工智能深度学习计算机视觉【免费下载链接】tensorrtxImplementation of popular deep learning networks with TensorRT network definition API项目地址https://gitcode.com/gh_mirrors/te/tensorrtx点击查看免费下载相关推荐tensorrtx 实战基于 TensorRT 网络定义 API 部署 ArcFace 人脸识别R50 / R100 / MobileFaceNettensorrtx 实战基于 TensorRT 网络定义 API 部署 ArcFace 人脸识别R50 / R100 / MobileFaceNet 本指人工智能深度学习计算机视觉ChineseOCR终极指南如何用YOLOv3CRNN实现中文场景文字识别ChineseOCR终极指南如何用YOLOv3CRNN实现中文场景文字识别 ChineseOCR是一个基于YOLOv3和CRNN深度学习技术的中文文字识别开计算机视觉深度学习后端QEMU Reset 机制深度解析Resettable 接口、三阶段复位与系统级复位流程QEMU Reset 机制深度解析Resettable 接口、三阶段复位与系统级复位流程 导读 QEMU 作为系统级模拟器需要精确模拟真实硬件中的各种复位人工智能深度学习计算机视觉上一篇Vue 3企业级组件库架构演进Element Plus如何解决大规模前端工程的技术债务问题下一篇The Wisdom of Quinn代码签名避坑清单6个Gatekeeper拦截与签名崩溃原因创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
阅读完成 · 觉得有帮助?