详细记录超轻量中文OCR LSTM模型ncnn实现

详细记录超轻量中文OCR LSTM模型ncnn实现

0x0 超轻量中文OCR检测识别一条龙项目

我也是contributor啦~

欧!耶!

https://github.com/ouyanghuiyu/chineseocr_litegithub.com

0x1 缘由

ncnn/src/layer/lstm.cpp 这份代码存在了非常久,我认为写的不对,一直是禁用编译的状态

只看LSTM前向的公式挺简单,自己没用过的东西必然是写不好的,就像以前写过一个tensorflow2ncnn.cpp 到处都是坑,自己又不会tensorflow于是也修不好

ouyanghuiyu的OCR一条龙项目用到了LSTM,效果不错,用作实验模型把LSTM搞出来提升识别准确性,一方面算是填上ncnn LSTM的坑,一方面也方便社区开发者更好的借(白)鉴(嫖)

ncnn没有batch维度,跑LSTM模型需要一些特殊手法,我把LSTM和OCR在ncnn上的实现过程写出来,作为参考

https://github.com/ouyanghuiyu/chineseocr_lite/pull/41github.com

0x2 导出onnx模型

git clone https://github.com/ouyanghuiyu/chineseocr_lite.git

整个项目中只有CRNN模型使用了LSTM,把单独调CRNN的部分写个单独的程序出来方便分析

CRNN模型输入的图片为不定长width x 32,灰度图 为了避免opencv PIL和ncnn在图片解码和resize上的差异,我准备了一张 277x32 尺寸的bmp灰度图片,作为测试图片

from config import  *
from crnn import FullCrnn,LiteCrnn,CRNNHandle
from PIL import Image
import numpy as np
import cv2

crnn_net =  LiteCrnn(32, 1, len(alphabet) + 1, nh, n_rnn=2, leakyRelu=False, lstmFlag=LSTMFLAG)

crnn_handle  =  CRNNHandle(crnn_model_path , crnn_net , gpu_id=0)

img = cv2.imread("testcrnn.bmp")

partImg = Image.fromarray(np.array(img))

partImg_ = partImg.convert('L')

simPred = crnn_handle.predict(partImg_)

print(simPred)

跑一下看看,输出是正常的

$ python testcrnn.py
device: cuda:0
load model
植物医生官方唯一授权

可以导出onnx模型了,找到模型加载的地方在 chineseocr_lite/crnn/CRNN.py init函数最后附加导出代码,再跑一次testcrnn.py就有onnx了

这里注意的是,onnx export函数需要加上keep_initializers_as_inputs=True, opset_version=11两个参数,否则后面onnx-simplifier可能会段错误,无法简化模型

output_onnx = 'crnn_lite_lstm_v2.onnx'
input_names = ["input"]
output_names = ["out"]
inputs = torch.randn(1, 1, 32, 277).to(self.device)
torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False, input_names=input_names, output_names=output_names, keep_initializers_as_inputs=True, opset_version=11)

直接导出的onnx模型有很多胶水op是ncnn不支持的,用onnx-simplifier是常规操作

pip install -U onnx --user
pip install -U onnxruntime --user
pip install -U onnx-simplifier --user

python -m onnxsim crnn_lite_lstm_v2.onnx crnn_lite_lstm_v2-sim.onnx

0x3 实现LSTM

其实原本的lstm.cpp大致完成度还是挺高的,稍微改改就可以了,比如加上forward reverse bidirectional三种方向,具体公式参考onnx LSTM的文档对着写就行

https://github.com/onnx/onnx/blob/master/docs/Operators.md#LSTMgithub.com

稍微注意的坑是onnx LSTM的权重layout,早期ncnn lstm.cpp的权重layout是跟着caffe做的,顺序是IFOG,而onnx的顺序是IOFG,为了保持兼容性,会在onnx2ncnn中自动转换。一开始没注意发现怎么老算不对...

onnx LSTM的定义中包含了许多训练阶段或者通常推理阶段不会用到的参数和输入,比如sequence_lens,初始的hidden/cell,可选的hidden/cell输出,这些不常用的东西ncnn中都去掉了。还有许多自定义的激活函数,我个人觉得大部分情况支持标准的LSTM sigmoid+tanh足够了,如果有新的以后再加不迟

batch维度的处理依旧遵循ncnn的convention直接砍掉,输入 [seqlen, batch, size] 变为 [seqlen, size],输出 [seqlen, batch, hidden] 变为 [seqlen, hidden]

https://github.com/Tencent/ncnn/pull/1613github.com

0x4 手工微调模型和ncnn推理实现

转模型是常规操作,这次没有任何报错,不报错不代表一定能用,先用netron工具打开param看看模型结构

$ ./onnx2ncnn crnn_lite_lstm_v2-sim.onnx crnn_lite_lstm_v2.param crnn_lite_lstm_v2.bin
Convolution      Conv_27                  1 1 103 104 0=512 1=1 11=1 2=1 12=1 3=1 13=1 4=0 14=0 15=0 16=0 5=1 6=262144
Squeeze          Squeeze_29               1 1 104 106 -23303=1,2
Permute          Transpose_30             1 1 106 107
LSTM             LSTM_37                  1 1 107 234 0=128 1=524288 2=2
Permute          Transpose_38             1 1 234 237 0=2
Reshape          Reshape_40               1 1 237 239 0=-1 1=0
Reshape          Reshape_54               1 1 239 253 0=256
InnerProduct     Gemm_55                  1 1 253 254 0=256 1=1 2=65536

看到这里红框和蓝框部分,有许多Squeeze/Reshape/Permute叠在一起,就要深刻怀疑模型结构大概率有缺陷,需要手工微调下才能使用,这里找找pytorch代码是何种操作,在chineseocr_lite/crnn/crnn_lite.py

# conv features
conv = self.cnn(input)
b, c, h, w = conv.size()

assert h == 1, "the height of conv must be 1"
conv = conv.squeeze(2)
conv = conv.permute(2, 0, 1)  # [w, b, c]

上面就是红框部分的代码,实际作用是把Convolution的输出4维blob,其中h一定是1,把h维砍掉reshape到3维,然后[b,c,w]换成[w,b,c]

ncnn没有batch维,所以在ncnn里可以等价为,Convolution的输出3维blob,其中h一定是1,把h维砍掉reshape到2维,然后[c,w]换成[w,c]。

我可以用Reshape和Permute实现这个过程。注意到Convolution的输出w是不定长的,因此Reshape的outw(0)参数用-1代表剩余长度表达不定长,-233表达3维降到2维;Permute层用0=1参数表达交换2维

Reshape          Squeeze_29               1 1 104 106 0=-1 1=512 2=-233
Permute          Transpose_30             1 1 106 107 0=1

关于手工修改模型结构,也可以参考这里的zhihu文章

nihui:手工优化ncnn模型结构zhuanlan.zhihu.com图标
recurrent, _ = self.rnn(input)
T, b, h = recurrent.size()
t_rec = recurrent.view(T * b, h)
output = self.embedding(t_rec)  # [T * b, nOut]

上面就是蓝框部分的代码,实际作用是把LSTM的输出3维blob,前面两维seqlen和batch合并reshape到2维,形成[seqlen * batch, hidden]送给后面做batch fc

ncnn没有batch维,因此这个reshape可以直接省略,但后面这个fc无论如何都要做batch,无法手工微调实现这个过程,便不调整

0x5 外部展开实现batch fc

虽然刚才蓝框的问题还没解决,整个模型结构看起来还不太对,但只要模型符合格式,加载依然是ok的,并且由于ncnn的partial inference的特点,只要你extract的路径不走到蓝框的部分,也是可以正常推理的

常规操作,ncnn加载模型,创建Extractor 我不直接拿output,我先拿LSTM的输出blob,然后手工循环seqlen次,就像遍历很多图片那样,每次引用一个hidden特征,创建新的子Extractor,算出这个hidden的输出,最后再拼回去就是batch fc的结果了

这种方法实现batch inference也是ncnn的推荐做法,ncnn example中的faster-rcnn例子也是这样实现了第二阶段多个proposal roi的分类

// lstm
ncnn::Mat blob234;
crnn_ex.extract("234", blob234);

int seqlen = blob234.h;

// batch fc
ncnn::Mat blob254(256, seqlen);
for (int i=0; i<seqlen; i++)
{
    ncnn::Extractor crnn_ex_1 = crnn_net.create_extractor();

    ncnn::Mat blob234_i = blob234.row_range(i, 1);
    crnn_ex_1.input("253", blob234_i);

    ncnn::Mat blob254_i;
    crnn_ex_1.extract("254", blob254_i);

    memcpy(blob254.row(i), blob254_i, 256 * sizeof(float));
}

0x5 总结

导出onnx,实现ncnn新的层,修改杂七杂八的胶水op,跳过部分模型的推理,实现batch,真是篇有实践性的文章,传递些经验

欢迎加qq群参与技术交流qwq 637093648

编辑于 03-15