【技术分享】⾯向隐私AI的TensorFlow深度定制化实践

作者：Rosetta技术团队

策划：蔡芳芳

在上⼀篇⽂章中， 【技术分享】隐私AI工程技术实践指南，我们整体上了介绍了基于深度学习框架开发隐私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形成的完整计算逻辑，在内部会以dat aflow graph的形式来表达。在这⼀有向⽆环图（DAG）上，以算⼦等作为节点，以Tesnor等作为边来指明数据的流动路径。在graph上，有些节点是TF框架⾃身根据需要添加的，⽐如，⽤户在training算法阶段时，只需要调⽤各种优化器（Optimizer）的minimize⽅法，TF⾃身就会⾃动找到前向图中各算⼦所对应的梯度算⼦，并按照数学上的链式求导法则，构建出反向梯度⼦图。

<img src="./figs/dag_nn.png"alt="TF执⾏过程示例"st yle="zoom:50%;"/>

 

TensorFlow数据流计算图（图⽚来⾃TensorFlow社区）

>>>>

Session（会话）

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

<img src="./figs/device_run.png"alt="TF分布式执⾏"st yle="zoom:30%;"/>

 

TensorFlow分布式并⾏执⾏（图⽚来⾃⽹络）

TensorFlow的codebase本身还是很复杂的，篇幅所限，难以在此对TensorFlow进⾏深⼊的介绍，感兴趣的读者可以参考InfoQ上其他优秀⽂章以进⼀步学习TensorFlow。

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

TF提供了⽐较丰富的扩展⽅法，除了在Python层可以基于内置的丰富算⼦集合，通过模块的继承、组装等⽅式得到⾃定义的功能之外，还可以在后端C++层⾃定义⾃⼰的算⼦[2]。在后端基于CustomC++op机制进⾏扩展相⽐于在前端层进⾏扩展有⼀些特别的优势：

• 有时候基于现有TF原⽣算⼦表达上层⾃定义逻辑很困难，⽽在后端实现则更灵活⾃由；

• 通过后端CustomC++op，可以以更加⾼效的⽅式实现⾃⼰的逻辑，可以在其中进⾏更底层的、⾯向编译器等的各种优化；

整体上看，基于TF的扩展⼯具，使⽤customC++op，只需要完成以下四步即可：

1. 通过TF提供的C++宏⼯具注册新的op。这主要是定义好这个op的输⼊输出类型、名称等接⼝信息。例如在Rosetta中可以如下定义⼀个新的op：

2. 在C++中具体的实现这个op所对应的内部处理逻辑，这就是所谓的后端“kernel”。TF提供了⼀些⽅便的基类接⼝，⽤户⼀般只需要定义⼀个⼦类，override实现其中的compute⽅法即可，例如：

3. 基于REGISTER_KERNEL_BUILDER这样的宏，将上⾯所定义的接⼝和内部的实现给绑定起来。这是因为TF⽀持基于不同的输⼊、输出类型和所运⾏的底层设备架构来定义同⼀个算⼦不同的内部实现，所以⽤户可以定义多种kernel实现，告知给系统什么场景下运⾏具体哪⼀个kernel，在实际运⾏时，TF就可以根据不同的设备、数据流上下⽂调⽤不同的kernel来实际执⾏此op。例如：

4. 将你的后端算⼦库编译为⼀个动态库so⽂件后，在Pyt hon层调⽤接⼝引⼊此模块，然后就可以如同调⽤原⽣算⼦⼀样的⽅式来调⽤这些⾃定义算⼦了。例如：

如果你需要在模型训练程序中调⽤这个⾃定义算⼦，你还需要在Pyt hon层通过@ops.RegisterGradient("XXXOp")来注册这个算⼦对应的梯度算⼦，通过这种⽅式，TF就可以在⾃动构建反向梯度图时⾃动的实现对⾃定义算⼦梯度的集成。

Rosetta利用TF这⼀扩展机制引⼊两类算⼦：中间过渡层RttOps算⼦库和隐私计算SecureOps算⼦库，前者是为了⽀持⾯向⾃定义数据类型的计算图的构建，后者是为了对接后端隐私计算功能，并在执⾏图时进⾏动态绑定。之所以从设计上区分这两类算⼦，是因为可以进⼀步解耦图的构建和图的执⾏，提供更多的灵活性。引⼊了这两个基础的算⼦库之后，就可以进⼀步的进⾏整体的改造了。

• RttOp算⼦库

  与后端MPC隐私计算完全⽆关的辅助中间层，⼀系列的“浮标”置位算⼦，⽀持⾃定义Tensor类型。其内部默认的实现逻辑是和对应的TF原⽣算⼦⼀样的。

  
  

• SecureOp算⼦库

  完整的前后端算⼦库，注册了对应的梯度函数；在内部实现中调⽤隐私协议层的抽象算⼦接⼝实现和TF的对接。

Rosetta 对 TensorFlow 的深度定制化

如上⼀篇⽂章整体介绍的那样，作为⾯向实际⼯业落地⽬标的隐私AI框架，Roset t a对于TF的改造原则始终是为了提供更加便于AI开发者使⽤的上层接⼝，以及兼顾系统后端隐私协议的可扩展性。

<img src="./figs/arch_det ail_cn.png"alt="Roset t a整体⼯程架构"style="zoom:40%;"/>

Rosetta整体⼯程架构

从系统架构和代码上看，改造的⼊⼝可以分为两⼤部分：

1. 后端C++部分的适配定制。主要以⾃定义算⼦的kernel形式进⾏适配。⼤部分接⼝的输⼊输出参数是以tf.string基础类型的Tensor，⾥⾯封装的是⾃定义的密⽂数据。在隐私算⼦SecureOps的kernel内部会进⼀步调⽤统⼀的密码协议接⼝来完成TF到隐私计算功能的联通。

   
   

2. 前端Python部分的适配定制。这⾥除了在Pyt hon前端引⼊我们⾃定义的算⼦库之外，还需要进⼀步改造TF中的⾃动求导功能等模块以实现对于新隐私算⼦的⾃动构建图、⾃动求导的⽀持。

从对程序的动态处理⻆度来看，如前⼀篇⽂章所说，Rosetta是经过两个阶段的Pass，来完成到底层多⽅协作的MPC处理程序的转换。这⾥⼤部分基于TF的前后端改造都是为了完成StaticPass阶段的转换，即将原⽣Tensor转换为⽀持⾃定义密⽂类型的RttTensor，将原⽣Operation转换为⽀持tf.string格式输⼊输出的RttOp ，并最终在图开始启动时进⼀步的转换为承载实际MPC操作的SecureOp 。

细⼼的读者可以看出，上⾯在介绍TF的customC++op扩展机制的同时，我们已经展示了如何定义Rosetta中的单个新算⼦。接下来，我们介绍⼀下如何基于这些算⼦实现计算图的分阶段转换。

计算图的转换构建过程

>>>>

引入 rosetta 库时

⽤户在前端执⾏import lattciex.rosetta之后，Rosetta就会⽤RttOp静态替换掉原⽣TF中对应的原⽣API算⼦，且各个原⽣Tensor也会被包装⼀层到RttTensor ，其与原⽣Tensor的主要区别是，其数据的基础类型是tf.string

，且对应的计算算⼦是RttOp。这种基础类型的转换是基于RttOp算⼦库中的TfToRtt和RttToTf两个⽤于类型转换的算⼦来完成的。

<img src="./figs/st at ic_replace.png"alt="Roset t a在import时的静态替换"style="zoom:50%;"/>

 

>>>>

调用 Session.run 接口时

我们同样hook了Session.run⼊⼝，在其内部完成从上⼀步骤中RttOp算⼦到SecureOp算⼦的转换。如果⽤户使⽤TensorBoard⼯具查看此时的运⾏图，就会看到我们在图上添加了⼀个和原⽣TF计算图基本同构的新⼦图，这个⼦图就是由SecureOp 构成。

<imgsrc="./figs/forward_secureop_dag.png"alt="TensorBoard可以查看得到的SecureOp计算图"style="zoom:50%;"/>

 

和上⽂介绍的原⽣TF中的完整图构建过程⼀样，如果⽤户的程序含有模型训练过程，调⽤了优化器Optimizer的minimize⽅法，则我们还需要完成对SecureOp的反向梯度图⾃动⽣成的⽀持。

⾸先，我们需要注册各个SecureOp算⼦所对应的梯度函数。⽐如对于隐私矩阵乘法算⼦SecureMatMul，我们按照底层梯度的计算逻辑，定义其梯度函数如下：

此外，由于我们使⽤tf.string来统⼀承载⾃定义的密⽂数据类型，⽽TF本身是不⽀持对于tf.string类型算⼦的⾃动求导的，所以Rosetta中还对tf.python.ops.qradients_util等⼊⼝进⾏了hook改造。⽐如，在下⾯这⾥，我们设定当tensor的基础类型为string时仍可以继续进⾏反向传播：

<img src="./figs/bp_hook.png"alt="反向梯度图的⾃动⽣成"st yle="zoom:90%;"/>

 

通过这些精细的定制化改造，最终就可以实现反向梯度⼦图的⾃动⽣成，可以极⼤的降低⽤户上⼿隐私计算的开发难度。

<img src="./figs/backward_secureop_replace.png"alt="反向梯度图的⾃动⽣成"style="zoom:80%;"/>

 

>>>>

补充说明

• 并⾮所有的算⼦都需要转换为SecureOp，这是因为如果⼀个局部⼦图中全部的输⼊都是本地的常量（公开的写定到代码中的数据，⽆需保护），那么就没有必要将这个⼦图转换为多⽅协作的隐私计算⽅式计算，这样可以减少不必要的计算时间。

• 转换时，由于此时知道了即将运⾏的完整⼦图的信息，⽐如DAG图上有多少了算⼦需要运⾏，所以可以在这⾥进⾏⼀些定制化的优化，⽐如优化底层协议中多⽅之间的并发通讯。

在通过上述过程完成在前端层到SecureOp图的构建后，接下⾥就是依赖TF⾃身的图执⾏引擎来调度执⾏各个SecureOp的后端实现了，在这个kernal中，为了和具体使⽤的隐私计算技术解耦，我们所调⽤的是密码协议接⼝，⽐如SecureMatMul⾥最终通过如下代码⽚段来调⽤内部“隐私计算引擎”。这⾥的内部细节，我们会在后续内容中加以介绍。

小结

在本篇⽂章中，我们进⼀步介绍了Rosetta是如何深度适配、定制化改造TensorFlow的各个组件以引⼊隐私计算功能的。与其他隐私AI开源框架相⽐，Rosetta由于需要同时对TensorFlow的前端和后端进⾏扩展，并且完全复⽤对上层的API接⼝，所以定制化的程度更加深⼊。这⾥的改造是偏向于“系统易⽤性”这⼀⽬标的，不需要太多涉及MPC等隐私计算技术，⾄于如何在后端引⼊”隐私计算引擎“，我们会在下⼀篇⽂章中介绍。

>>>>

作者介绍

Rosetta技术团队，⼀群专注于技术、玩转算法、追求⾼效的⼯程师。Rosetta是⼀款基于主流深度学习框架TensorFlow的隐私AI框架，作为矩阵元公司⼤规模商业落地的重要引擎，它承载和结合了隐私计算、区块链和AI三种典型技术。

⽬前Rosetta已经在Github开源（https://github.com/LatticeX-Foundation/Rosetta），欢迎关注并参与到Rosetta社区中来。

> >>>

参考文献

[1]Abadi,Martín,etal."Tensorflow:A system for large-scale machine learning."*12t h{USENIX}symposium on operat ing syst ems design and implement at ion({OSDI}16)*.2016.

[2] TensorFlow对定制化Op扩展的⽀持：https://www.tensorflow.org/guide/create_op

本文分享自微信公众号 - GDG（GDG_Shanghai）。
如有侵权，请联系 support@oschina.cn 删除。
本文参与“OSC源创计划”，欢迎正在阅读的你也加入，一起分享。

来源：oschina

链接：https://my.oschina.net/u/3989964/blog/4647617