面向隱私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");

2. 在 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...
}
}

3. 基於REGISTER_KERNEL_BUILDER這樣的宏，將上面所定義的接口和內部的實現給綁定起來。這是因爲 TF 支持基於不同的輸入、輸出類型和所運行的底層設備架構來定義同一個算子不同的內部實現，所以用戶可以定義多種kernel實現，告知給系統什麼場景下運行具體哪一個kernel，在實際運行時，TF 就可以根據不同的設備、數據流上下文調用不同的kernel來實際執行此 op。例如：

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

4. 將你的後端算子庫編譯爲一個動態庫 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 工程技術實踐指南：整體介紹