使用 TensorRT 网络定义 API 部署 CRNN 场景文字识别：tensorrtx/crnn 实战指南

导读

本文以开源仓库 tensorrtx 中的 crnn 模块为主体，完整讲解如何将 PyTorch 训练好的 CRNN（卷积循环神经网络）场景文字识别模型一步步转换为 TensorRT 引擎并部署推理：从权重导出（wts 文件）、CMake 构建，到纯 TensorRT Network Definition API 的 CNN + 双向 LSTM 网络搭建，再到 CTC 解码与结果验证。读完本文，你将掌握一套无需解析器、完全基于 TensorRT 原生 API 手写网络的部署方法论，可直接复用于其他 RNN/LSTM 类模型。

---

一、CRNN 模型与 tensorrtx/crnn 的实现定位

CRNN（Convolutional Recurrent Neural Network）是经典的场景文字识别（STR）方案，结构上由 CNN 卷积主干、RNN（循环层）和 CTC 转录层三部分组成，能够直接对整行不定长文本图像做序列识别。本项目（tensorrtx，即 "Implementation of popular deep learning networks with TensorRT network definition API"）采用**网络定义 API（Network Definition API）**而非 ONNX/TensorFlow 解析器来重建 CRNN，这意味着网络中的每一个卷积、池化、LSTM 单元都是通过 nvinfer1 接口逐层手动构建的。

crnn 模块的核心文件构成：

文件	作用
crnn/crnn.cpp	网络定义、引擎序列化/反序列化、推理与 CTC 解码主程序
crnn/genwts.py	将 PyTorch 权重 crnn.pth 导出为 TensorRT 可读的 crnn.wts
crnn/CMakeLists.txt	构建脚本，链接 CUDA / OpenCV / TensorRT 库
crnn/logging.h	TensorRT 日志器封装（仓库内多模块共用）
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.txt）

crnn/CMakeLists.txt 明确了运行环境要求，这也是让上面两条命令能跑通的前提：

• C++ 标准：add_definitions(-std=c++11) 与 set(CMAKE_CXX_STANDARD 11)，并以 Debug 类型构建；
• CUDA 依赖：find_package(CUDA REQUIRED)，链接 cudart（CUDA 运行时）；
• TensorRT 依赖：链接 nvinfer；
• OpenCV 依赖：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.wts：

f.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、类别数 37（36 个字符 + 1 个 blank）以及 LSTM 隐藏单元数 256——这些数字必须与 C++ 端 crnn/crnn.cpp 中的网络定义严格一致，否则加载权重时会出现张量尺寸不匹配。

4.2 wts 文件格式

crnn.wts 是一个纯文本权重文件，格式与 tensorrtx 其他模块统一，可用 crnn/crnn.cpp 中的 loadWeights 函数印证：

• 首行：一个整数，表示权重张量的总数（即 state_dict 的键数量）；
• 其后每行：[权重名] [元素个数] [以空格分隔的大端序 float 十六进制数据...]。

C++ 侧读取时逐行解析，按名存入 std::map<std::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 层卷积 + 批归一化 + 池化

createEngine（crnn/crnn.cpp）首先注册输入张量：

ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});

即单通道 32×100 的灰度图像，数据布局为 CHW。随后按 CRNN 经典结构逐层堆叠，关键调用的顺序为：

1. convRelu(..., 0) → 最大池化（2×2，stride 2）；
2. convRelu(..., 1) → 最大池化（2×2，stride 2）；
3. convRelu(..., 2, true)（带 BatchNorm）→ convRelu(..., 3) → 最大池化（2×2，stride (2,1)、padding (0,1)）；
4. convRelu(..., 4, true) → convRelu(..., 5) → 最大池化（2×2，stride (2,1)、padding (0,1)）；
5. convRelu(..., 6, true) 收尾。

convRelu 的实现（crnn/crnn.cpp）用 addConvolutionNd 建卷积层，按 ks/ss/ps/nm 数组设置卷积核尺寸、步长、padding 和输出通道，卷积后接 addActivation(..., ActivationType::kRELU) 做 ReLU 激活。第 2、4、6 层额外调用 addBatchNorm2d 挂接 BatchNorm。

BatchNorm 折叠技巧：TensorRT 的 addBatchNorm2d（crnn/crnn.cpp）并没有直接提供 BN 层，而是利用公式将 BN 折叠进 IScaleLayer（per-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/.bias，mean/var 来自 .running_mean/.running_var，eps 取 1e-5。这一手法在 tensorrtx 的众多视觉模型中通用，理解它有助于阅读仓库内其他模块（如 resnet/resnet50.cpp）的 BN 处理。

5.3 时序维度变换与双向 LSTM

CNN 输出经过 Shuffle 转置 Permutation{1, 2, 0} 后，进入两层双向 LSTM（crnn/crnn.cpp）：

auto lstm0 = addLSTM(network, weightMap, *sfl->getOutput(0), 256, "rnn.0.rnn");
// reshape 为 {26, 1, 1, 512} 后接 FC（512 → 256，embedding 层）
// 再转置 reshape 回 {1, 26, 256}，输入第二层 LSTM
auto lstm1 = addLSTM(network, weightMap, *sfl->getOutput(0), 256, "rnn.1.rnn");
// reshape 为 {26, 1, 1, 512} 后接 FC（512 → 37），作为最终输出

addLSTM（crnn/crnn.cpp）基于 TensorRT 的 addRNNv2 实现，要点有二：

• 方向：lstm->setDirection(RNNDirection::kBIDIRECTION)，即双向 LSTM，每层输出维度为隐藏单元数 × 2（256 → 512）；
• 权重拆分：PyTorch 的 LSTM 权重是 4 个门（input/forget/cell/output）拼接在一起的整体张量，因此 splitLstmWeights（crnn/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.cpp）：

builder->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.png（README 要求从 crnn.pytorch/data 复制到 crnn 目录），处理流程与 CRNN 训练约定一致（crnn/crnn.cpp）：

cv::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.at<uchar>(i) / 255.0 - 0.5) * 2.0;  // 归一化到 [-1, 1]

如果 demo.png 缺失，程序会打印 demo.png not found !!! 并退出。

6.3 异步推理

doInference（crnn/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，再交给 strDecode（crnn/crnn.cpp）解码：

std::string strDecode(std::vector<int>& 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]]);                  // sim：CTC 折叠
        }
    }
}

• raw 模式：不做任何处理，直接把 26 个 argmax 索引映射为字母。由于 blank 符 - 在索引 0 处，raw 输出会包含大量 -，例如 a-----v--a-i-l-a-bb-l-e---；
• sim 模式：执行标准的 CTC collapse 两步操作——① 跳过索引 0（blank 占位符）；② 合并相邻的连续重复字符。于是 raw 序列 a-----v--a-i-l-a-bb-l-e--- 折叠为 available。

这正是 CTC 损失训练模型的标准解码行为，也是校验部署正确性的最直接手段：只要 sim 输出等于预期文本，就说明权重加载、网络构建与推理链路完全正确。

---

八、工程要点与常见问题

1. 权重路径：createEngine 中写死从 ../crnn.wts 加载权重（crnn/crnn.cpp），因此按 README 流程应将 crnn.wts 放入 tensorrtx/crnn 目录，并从 build/ 子目录发起构建，才能解析到上一级的 wts 文件；
2. 图片路径：推理阶段读取的是当前工作目录下的 demo.png（crnn/crnn.cpp），务必在 build/ 中先复制 demo.png 再执行 ./crnn -d；
3. 精度开关：USE_FP16 宏控制 FP16/FP32；FP16 可显著降低带宽与加速，但需要硬件支持（Tesla/RTX 等），不支持时会退化为不可用，需关闭该宏；
4. BatchSize 固定为 1：BATCH_SIZE、输入输出缓冲区尺寸均在源码中硬编码，如需批处理需同步修改 main 中 static float data[...] 等缓冲区大小及 OUTPUT_SIZE 相关计算；
5. OCR 文本为小写字母与数字：模型字母表仅含 -0123456789abcdefghijklmnopqrstuvwxyz，输入文本需为小写，超出字母表范围的字符不会被正确识别；
6. 输入宽高约定：网络输入固定为 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 推理实现与示例验证，不包含完整的文本纠错与业务后处理，这部分需要依据实际业务自行扩展。