面向隐私计算+AI的TensorFlow深度定制化实践

面向隐私AI的TensorFlow深度定制化实践

在上一篇文章整体介绍中，我们整体上了介绍了基于深度学习框架开发隐私 AI 框架中的工程挑战和可行解决方案。在这一篇文章中，我们进一步结合 Rosetta 介绍如何定制化改造 TensorFlow 中前后端相关组件，以集成 MPC 等隐私计算技术，同时保留对 TensorFlow 接口 API 的复用，从而实现我们上一篇文章中所强调的“系统易用性”。

目前 Rosetta 主要基于 TensorFlow 1.14 CPU 版本加以开发（以下简称 TensorFlow 为 TF），这是因为 TF 1.x 目前在工业界中实际应用是较为广泛的，而引入动态图等高级功能的 TF 2.0，则由于接口的不向后兼容等问题，仍没有得到大规模的落地。 后续我们也将在 Rosetta 本身功能稳定的基础上考虑支持 TF 2.0。下面就让我们开始吧。

TensorFlow 快速回顾

想要基于 AI 框架进一步扩展引入隐私计算功能，第一步就是需要比较深入的了解这些 AI 框架，所以首先让我们简单回顾一下TF中核心的一些概念以及其宏观的内部处理过程。

TensorFlow 的核心概念

• Tensor（张量）

深度学习需要完成对大量高维度复杂数据的处理，在TensorFlow中，用Tensor来封装同一类型数据的高维数组。其中，基础类型除了各种不同精度的整数、浮点数外，还支持tf.string类型，这给我们提供了进行自定义类型改造的可能性。

一个三维Tensor (图片来自网络)

• Operation（算子）

Operation（算子，有时也称“操作”）用来封装对于 Tensor 的处理逻辑。同时也是连接 TF 的前端和后端之间逻辑处理的基本单元，在实际使用中，用户可以使用keras等上层封装 API 更方便的表达复杂计算逻辑，但是这些上层模块的内部，也会最终调用各个算子来完成逻辑的表达。

• Graph（计算图）

用户在 TF 前端调用各 API 形成的完整计算逻辑，在内部会以 dataflow graph 的形式来表达。在这一有向无环图（DAG）上，以算子等作为节点，以 Tesnor 等作为边来指明数据的流动路径。在 graph 上，有些节点是 TF 框架自身根据需要添加的，比如，用户在training算法阶段时，只需要调用各种优化器（Optimizer）的minimize方法，TF 自身就会自动的找到前向图中各算子所对应的梯度算子，并按照数学上的链式求导法则，构建出反向梯度子图。

​ TensorFlow 数据流计算图 (图片来自 TensorFlow 社区)

• Session（会话）

Session 主要是在实际执行 graph 时对一次执行的上下文进行维护处理。当用户调用其run方法时，TF 就会分析为了获取这一次的计算目标所需要运行的子图，并结合 TF 内置的强大的并行优化、分布式执行等模块，将所需要执行的逻辑进一步拆分为各个子图，各自映射到当前的可用设备资源上，最终调度这些设备以并行的方式高效完成计算任务。

​ TensorFlow 分布式并行执行 (图片来自网络)

TensorFlow 的 codebase 本身还是很复杂的，篇幅所限，难以在此对 TensorFlow 进行深入的介绍，感兴趣的读者可以参考 TensorFlow 社区中其他优秀文章以进一步学习。

TensorFlow 自定义算子库的扩展方法

TF 提供了较为丰富的扩展方法，除了在 Python 层可以基于内置的丰富算子集合，通过模块的继承、组装等方式得到自定义的功能之外，还可以在后端 C++ 层自定义自己的算子 [2]。在后端基于 Custom C++ op 机制进行扩展相比于在前端层进行扩展有一些特别的优势：

• 有时候基于现有 TF 原生算子表达上层自定义逻辑很困难，而在后端实现则更灵活自由；
• 通过后端 Custom C++ op，可以以更加高效的方式实现自己的逻辑，可以在其中进行更底层的、面向编译器等的各种优化；

整体上看，基于 TF 的扩展工具࿰