爱作业口算同步练，TensorFlow Lite 实践
文 / 杭州大拿科技股份有限公司
爱作业是一款通过 AI 技术自动批改数学作业的 APP，上线一年来已经积累超过千万用户， 帮助大量的老师和家长节约时间，提高效率。
在爱作业 APP 里面有个功能叫 “口算同步练”，我们希望提供给小朋友一种轻量的方便的方法可以直接在手机上巩固训练口算。如下图所示：
识别在屏幕上手写的笔迹，我们最早使用的方法是追踪书写轨迹，这也是比较传统和经典的方法，这种方法依赖于规则，不需要大量的训练数据。 但是经过我们试验，这种方法对于少量数字效果还不错，对于位数比较多的数字，特别是数字之间还存在连写，或者书写顺序不够标准的情况， 识别效果就差强人意了。
这时候我们想到要使用 Deep Learning 来解决问题，然后我们需要解决训练数据，模型， 和移动端部署的问题。
首先是训练数据，这是深度学习绕不开的问题。为此，我们专门开发了一个小工具（iOS / Android），让我们的标注人员可以方便的生成训练数据。这个工具会随机生成一些数字（包含小数点，分数），然后标注人员在屏幕中书写提示的数字，然后可以方便的保存到我们的服务器。 我们就这样收集到了第一批原始数据大概有 5000 张训练图片左右。
然后是模型。由于模型需要在移动端运行，所以不管是图片输入大小和模型本身，都需要相对精简。基于竖屏的比例，我们最终采用了将所有图片预处理成 60x100 (高x宽) 的大小。模型的整体方案是一个 seq2seq 的模型，这个框架最早是 Google 用于解决 NLP 相关的问题， 我们可以把它做一些调整来做图像 OCR。
首先我们需要一个生成 feature 的基础模型，最终我们采用了 MobileNet。 MobileNet 模型的核心就是将原本标准的卷积操作因式分解成一个 depthwise convolution 和一个 1*1 的卷积（文中叫 pointwise convolution）操作。简单讲就是将原来一个卷积层分成两个卷积层，其中前面一个卷积层的每个 filter 都只跟 input 的每个 channel 进行卷积，然后后面一个卷积层则负责 combining，即将上一层卷积的结果进行合并。MobileNet 可以在基本保证准确率的前提下大大减少计算时间和参数数量。 正如这个命名，我们相信用这个作为基础网络可以在移动端有不错的表现， 并且这个模型在 TensorFlow 的 model zoo 中有标准实现，并且提供了预训练模型。丰富的模型实现也是我们当初选择用 TensorFlow 来实现的重要原因之一。
在得到图片 feature 以后，我们通过 LSTM RNN 进行一个标准的 Encoder->Decoder 过程，得益于 TensorFlow 的丰富实现，这整个过程只要调用 TensorFlow 中 seq2seq 相关接口就可以轻松实现。 不过 seq2seq 的接口在 TensorFlow 1.0 以后有一些接口上的变化，我们使用了更新以后的接口。
接下来是如何把模型放到客户端去运行，我们在这里面临诸多选择。 整体方案来看，我们需要选择使用 TensorFlow Mobile 还是 TensorFlow Lite，另外在 iOS 端，我们需要考虑是否要将模型转为 iOS 的 CoreML 模型。由于 CoreML 只从 iOS 11 开始支持，我们很快否掉了这个选项。我们在做这个功能的时候，TensorFlow 刚发布了 1.4（目前已经 1.12 了），Lite 还是一个全新的东西， 最大的问题是对于各种 ops 支持不够好，而且当时官方网站也是推荐在生产环境中使用 Mobile 版本， 所以尽管 Lite 有着各种优点， 我们还是选择了 Mobile。
TensorFlow Mobile 的文档也算是相对丰富的，我们在部署的过程中并没有遇到太大的障碍。有一点需要注意的是，需要使用 print_selective_registration_header 这个小工具来分析模型中真正使用到的算子，从而大大减少生成的运行库的大小。
就这样我们第一个版本顺利上线并稳定运行，但我们依然关注着 Lite 的动态。 在经过一些考虑后，特别是在 2018 年 10 月的谷歌开发者大会上听到了更多 TensorFlow Lite 的介绍，我们决定用 Lite 替换掉 Mobile 版本，主要基于以下几个考虑：
更小的运行库体积， 通过 TensorFlow Lite 生成的运行库体积很小，去掉了很多不必要的依赖。比如模型采用新的更轻量级 FlatBuffers 格式，而之前使用的 protobuf 需要很多的依赖库。同时 TensorFlow Lite 以更加轻量级的方式实现了原来 TensorFlow 中的核心算子， 也有效降低了运行库的体积。
更快的运行速度。 TensorFlow Lite 很多核心算子是专门为移动平台优化过的，同时通过 NNAPI 能够透明的支持 GPU 加速，在一些老设备上能够自动回滚到 CPU 计算。
TensorFlow Lite 的模型量化功能，将模型量化到 int8，在几乎不怎么损失计算精度的情况下把模型变为原来的 1/4 大小。
但是我们仍然遇到了很多困难。主要是：
TensorFlow Lite 尚不支持 control flow，意味着 TensorFlow 新 seq2seq 接口中的 dynamic_decode 不能正确导出执行。
TensorFlow Lite 中对于 LSTM 的支持不够完善，比如 forget bias 参数并不支持。
我们所使用的部分 operators 在 TensorFlow Lite 中还不支持，比如 GatherTree等等。
在 TensorFlow Lite 中实现 control flow 并不是件容易的事情，并不能直接通过 custom op 来实现，最终我们把动态解码变成了静态展开，稍微损失了一些性能但是解决了这个问题。 对于 forget bias 的问题， 我们取巧的在参数恢复的时候把 forget bias 加在 LSTM forget gate 的 bias 上，这样就不用修改 TensorFlow Lite 的代码了。对于其他问题，都可以通过 custom op 来实现，相对容易一些。
最终我们通过 toco 把 graph pb 导出成 TensorFlow Lite 格式，接下来就是要集成到手机上运行。在这一步上我们并没有遇到太多阻碍。因为有了之前集成 TensorFlow Mobile 的经验，我们很快完成了集成。 我们在具体实现的时候把识别接口做了封装，这样在我们从 TensorFlow Mobile 迁移到 TensorFlow Lite 的时候，客户端开发的同事几乎不需要改动代码就可以使用新的接口。
在完成 TensorFlow Mobile 到 TensorFlow Lite 的迁移以后，我们单次 inference 的速度提高了 20%，模型大小减少了 75%，可以说是非常有吸引力的。
随着手机端计算能力的不断增强，越来越多的深度学习模型可以在手机端运行。 目前 TensorFlow Lite 也在快速开发迭代中，相信 TensorFlow Lite 会在更多实践中越来越好。
感谢 TensorFlow Lite 团队在我们的开发过程中给予的支持，特别是顾仁民和刘仁杰两位技术专家给我们提供了很多有帮助的资料和参考，希望将来有更多的合作机会。