目前Rosetta主要基于TensorFlow 1.14 CPU版本加以开发(以下简称TensorFlow为TF),这是因为TF 1.x 目前在工业界中实际应用较为泛,而引入动态图等级功能的TF 2.0,则由于接口不向后兼容等问题,仍没有得到大规模落地。后续我们也将在Rosetta本身功能稳定的基础上考虑支持TF2.0。下面就让我们开始吧
1、TensorFlow快速回顾
想要基于AI框架进步扩展引入隐私计算功能,第步需要比较深入地了解这些AI框架,所以首先让我们简单回顾下TF的核心概念以及宏观的内部处理过程
TensorFlow核心概念
Tensor(张量)深度学习需要完成对大量高维度复杂数据的处理,在TensorFlow中,用Tensor来封装同 类型数据的维数组。其中,基础类型除了各种不同精度的整数、浮点数外,还持tf.string类型,这给我们提供了进行自定义类型改造的可能性。
个三维Tensor|图来络
Operation(算子) Operation(算子,有时也称“操作”)用来封装对于 Tensor 的处理逻辑。同时也是连接TF的前端和后端之间逻辑处理的基本单元,在实际使用中,用户可以使用keras等上层封装API更便的表达复杂计算逻辑,但是这些上层模块的内部,最终也会调用各个算来完成逻辑的表达。
Graph(计算图)用户在TF前端调用各API形成的完整计算逻辑,在内部会以dataow 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的扩展工具,使用custom C++ op,只需要完成以下四步即可:
1. 通过TF提供的 C++ 宏工具注册新的op。这主要是定义好这个op的输入输出类型、名称等接口信息。例如在Rosetta中可以如下定义个新的op:
2. 在C++中具体的实
2. 在C++中具体的实现这个op所对应的内部处理逻辑,这就是所谓的后端 “kernel”。TF提供了些方便的基类接口,用户般只需要定义个子类,override实现其中的compute 法即可,例如:
3. 基于REGISTER_KERNEL_BUILDER这样的宏,将上面所定义的接口和内部的实现给绑定起来。这是因为TF支持基于不同的输入、输出类型和所运的底层设备架构来定义同个算子不同的内部实现,所以用户可以定义多种 kernel 实现,告知给系统什么场景下运行具体哪个 kernel ,在实际运行时,TF就可以根据不同的设备、数据流上下文调用不同的 kernel 来实际执此op。例如:
4. 将你的后端算子库编译为个动态库so文件后,在Python层调用接口引入此模块,然后就可以如同调用原生算子样的方式来调用这些自定义算子了。例如:
如果你需要在模型训练程序中调用这个自定义算子,你还需要在Python层通过@ops.RegisterGradient("XXXOp")来注册这个算子对应的梯度算子,通过这种式,TF就可以在动构建反向梯度图时自动的实现对自定义算子梯度的集成。
Rosetta利用TF这扩展机制引两类算子:中间过渡层 RttOps 算子库和隐私计算SecureOps算子库,前者是为了支持面向自定义数据类型的计算图的构建,后者是为了对接后端隐私计算功能,并在执行图时进行动态绑定。
之所以从设计上区分这两类算子,是因为可以进步解耦图的构建和图的执行,提供更多的灵活性。引入了这两个基础的算子库之后,就可以进步地进行整体的改造了。
* RttOp算子库
与后端MPC隐私计算完全无关的辅助中间层,系列的“浮标”置位算子,支持自定义Tensor类型。其内部默认的实现逻辑是和对应的TF原生算子样的。
* SecureOp算子库
完整的前后端算子库,注册了对应的梯度函数;在内部实现中调用隐私协议层的抽象算子接口实现和TF的对接。
2、Rosetta对TensorFlow的深度定制化
如上篇文章整体介绍的那样,作为面向实际工业落地目标的隐私AI框架,Rosetta对于TF的改造原则始终是为了提供更加便于AI开发者使用的上层接,以及兼顾系统后端隐私协议的可扩展性。
Rosetta整体程架构
从系统架构和代码上看,改造的入口可以分为两大部分:
1. 后端C++部分的适配定制。主要以自定义算子的 kernel 形式进行适配。大部分接的输入输出参数是以 tf.string 基础类型的 Tensor ,里面封装的是自定义的密文数据。在隐私算子SecureOps的 kernel 内部会进一步用统的密码协议接口来完成TF到隐私计算功能的联通。
2. 前端Python部分的适配定制。这里除了在Python前端引我们定义的算子库之外,还需要进步改造TF中的自动求导功能等模块以实现对于新隐私算子的自动构建图、自动求导的持。从对程序的动态处理角度来看,如前篇文章所说,Rosetta是经过两个阶段的Pass,来完成到底层多协作的MPC处理程序的转换。这里大部分基于TF的前后端改造都是为了完成Static Pas阶段的转换,即将原生 Tensor 转换为持定义密文类型的 RttTensor,将原生 Operation 转换为持 tf.string 格式输入输出的 RttOp,并最终在图开始启动时进步 的转换为承载实际MPC操作的执行命令:SecureOp。
细的读者可以看出,上面在介绍TF的custom C++ op扩展机制的同时,我们已经展示了如何定义Rosetta中的单个新算子。接下来,我们介绍下如何基于这些算子实现计算图的分阶段转换。
计算图的转换构建过程
引入Rosetta库时用户在前端执 import lattciex.rosetta 之后Rosetta就会用RttOp静态替换掉原生TF中对应的原生 API 算子,且各个原生Tensor也会被包装层到 RttTensor ,其与原生Tensor的主要区别是,其数据的基础类型是 tf.string ,且对应的计算算子是 RttOp。这种基础类型的转换是基于RttOp算库中的 TfToRtt 和 RttToTf 两个用于类型转换的算子来完成的。
调用Session.run接口时我们同样hook了 Session.run ,在其内部完成从上步骤中 RttOp 算子到SecureOp算子的转换。如果用户使用TensorBoard具查看此时的运图,就会看到我们在图上添加了个和原生TF计算图基本同构的新子图,这个子图就是SecureOp构成。
和上文介绍的原生TF中的完整图构建过程样,如果用户的程序含有模型训练过程,调用了优化器Optimizer的 minimize 法,则我们还需要完成对 SecureOp的反向梯度图动成的持。
首先,我们需要注册各个 SecureOp 算所对应的梯度函数。比如对于隐私矩阵乘法算子 SecureMatMul ,我们按照底层梯度的计算逻辑,定义其梯度函数如下:
此外,由于我们使用 tf.string 来统承载定义的密文数据类型,而TF本身是不支持对于 tf.string 类型算子的动求导的,所以Rosetta中还对执行命令: tf.python.ops.gradients_util 等入口进行了 hook 改造。比如,在下面这里,我们设定当tensor的基础类型为string时仍可以继续进反向传播:
通过这些精细的定制化改造,最终就可以实现反向梯度子图的自动生成,可以极大的降低用户上手隐私计算的开发难度。
补充说明
并非所有的算子都需要转SecureOp,这是因为如果个局部子图中全部的输入都是本地的常量(公开的写定到代码中的数据,无需保护),那么就没有必要将这个子图转换为多方协作的隐私计算式计算,这样可以减少不必要的计算时间。
转换时,由于此时知道了即将运的完整子图的信息,如DAG图上有多少了算子需要运,所以可以在这里进行些定制化的优化,比如优化底层协议中多方之间的并发通讯。在通过上述过程完成在前端层到 SecureOp 图的构建后,接下里就是依赖TF自身的图执引擎来调度执各个SecureOp的后端kernel实现了,在这个kernel中,为了和具体使用的隐私计算技术解耦,我们所调用的是密码协议接,如SecureMatMul最终通过如下代码片段来调用内部“隐私计算引擎”。这里的内部细节,我们会在后续内容中加以介绍。
3、小结
在本篇章中,我们进步介绍了Rosetta是如何深度适配、定制化改造 TensorFlow 的各个组件以引入隐私计算功能的。与其他隐私AI开源框架相,Rosetta由于需要同时对TensorFlow的前端和后端进行扩展,并且完全复用对上层的 API 接,所以定制化的程度更加深。这里的改造是偏向于“系统易用性”这一目标的,不需要太多涉及 MPC 等隐私计算技术,至于如何在后端引入”隐私计算引擎“,我们会在下篇文章中介绍。
-----作者介绍-----
Rosetta技术团队,群专注于技术、玩转算法、追求高效的工程师。Rosetta是款基于主流深度学习框架TensorFlow 的隐私AI框架,作为矩阵元公司大规模商业落地的重要引擎,它承载和结合了隐私计算、区块链和AI三种典型技术。目前Rosetta已经在Github开源:欢迎关注并参与到Rosetta社区中来!
举报/反馈

量子位

139万获赞 29.1万粉丝
追踪AI技术和产品新动态
北京极客伙伴科技有限公司
关注
0
0
收藏
分享