面向隐私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 进行深入的介绍，感兴趣的读者可以参考 InfoQ 上其他优秀文章以进一步学习 TensorFlow。

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

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

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

整体上看，基于 TF 的扩展工具，使用 custom C++ op，只需要完成以下四步即可：

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

REGISTER_OP("RttMatmul")
.Input("x: string")
.Input("y: string")
.Output("res: string")
.Attr("transpose_a: bool = false")
.Attr("transpose_b: bool = false");

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

template <typename Device>
class RttMatMulOp : public OpKernel {
public:
explicit RttMatMulOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("transpose_a", &transpose_a_));
OP_REQUIRES_OK(context, context->GetAttr("transpose_b", &transpose_b_));
}

void Compute(OpKernelContext* context) override {
// Check if the dimensions of the two matrices are valid
const Tensor& x = context->input(0);
const Tensor& y = context->input(1);
// detailed implementation...
}
}

1. 基于REGISTER_KERNEL_BUILDER这样的宏，将上面所定义的接口和内部的实现给绑定起来。这是因为 TF 支持基于不同的输入、输出类型和所运行的底层设备架构来定义同一个算子不同的内部实现，所以用户可以定义多种kernel实现，告知给系统什么场景下运行具体哪一个kernel，在实际运行时，TF 就可以根据不同的设备、数据流上下文调用不同的kernel来实际执行此 op。例如：

REGISTER_KERNEL_BUILDER(Name("RttMatmul").Device(DEVICE_CPU), RttMatMulOp<CPUDevice>);

1. 将你的后端算子库编译为一个动态库 so 文件后，在 Python 层调用接口引入此模块，然后就可以如同调用原生算子一样的方式来调用这些自定义算子了。例如：

# load librtt_ops.so
_rtt_ops_lib = os.path.dirname(__file__) + '/../../../librtt-ops.so'
rtt_ops = tf.load_op_library(_rtt_ops_lib)
# now, you can use the ops in this library as rtt_ops.rtt_matmul

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

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

• RttOp 算子库 与后端 MPC 隐私计算完全无关的辅助中间层，一系列的“浮标”置位算子，支持自定义 Tensor 类型。其内部默认的实现逻辑是和对应的 TF 原生算子一样的。
• SecureOp 算子库 完整的前后端算子库，注册了对应的梯度函数；在内部实现中调用隐私协议层的抽象算子接口实现和 TF 的对接。

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 Pass 阶段的转换，即将原生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，我们按照底层梯度的计算逻辑，定义其梯度函数如下：

@ops.RegisterGradient("SecureMatmul")
def SecureMatMulGrad(op, grad):
"""The gradient for the Secure MatMul operator."""
t_a = op.get_attr("transpose_a")
t_b = op.get_attr("transpose_b")
a = op.inputs[0]
b = op.inputs[1]
if not t_a and not t_b:
grad_a = SecureMatMul(grad, b, transpose_b=True)
grad_b = SecureMatMul(a, grad, transpose_a=True)
elif not t_a and t_b:
grad_a = SecureMatMul(grad, b)
grad_b = SecureMatMul(grad, a, transpose_a=True)
elif t_a and not t_b:
grad_a = SecureMatMul(b, grad, transpose_b=True)
grad_b = SecureMatMul(a, grad)
elif t_a and t_b:
grad_a = SecureMatMul(b, grad, transpose_a=True, transpose_b=True)
grad_b = SecureMatMul(grad, a, transpose_a=True, transpose_b=True)
return grad_a, grad_b

此外，由于我们使用tf.string来统一承载自定义的密文数据类型，而 TF 本身是不支持对于tf.string类型算子的自动求导的，所以 Rosetta 中还对tf.python.ops.gradients_util等入口进行了 hook 改造。比如，在下面这里，我们设定当 tensor 的基础类型为 string 时仍可以继续进行反向传播：

通过这些精细的定制化改造，最终就可以实现反向梯度子图的自动生成，可以极大的降低用户上手隐私计算的开发难度。

补充说明

• 并非所有的算子都需要转换为SecureOp，这是因为如果一个局部子图中全部的输入都是本地的常量（公开的写定到代码中的数据，无需保护），那么就没有必要将这个子图转换为多方协作的隐私计算方式计算，这样可以减少不必要的计算时间。
• 转换时，由于此时知道了即将运行的完整子图的信息，比如 DAG 图上有多少了算子需要运行，所以可以在这里进行一些定制化的优化，比如优化底层协议中多方之间的并发通讯。

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

// call protocol ops
vector<string> outstr(m*n);
ProtocolManager::Instance()->GetProtocol()->GetOps(msg_id().str())->Matmul(in1, in2, outstr, &attrs_);

小结

在本篇文章中，我们进一步介绍了 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, et al. “Tensorflow: A system for large-scale machine learning.” 12th {USENIX} symposium on operating systems design and implementation ({OSDI} 16). 2016.

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

系列文章：

隐私 AI 工程技术实践指南：整体介绍