Categories
程式開發

面向隱私AI的TensorFlow深度定制化實踐


在上一篇文章中,我們整體上了介紹了基於深度學習框架開發隱私AI 框架的工程挑戰和可行解決方案。 在這一篇文章中,我們進一步結合 羅塞塔 介紹如何定制化改造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類型,這給我們提供了進行自定義類型改造的可能性。

面向隱私AI的TensorFlow深度定制化實踐 1

一個三維Tensor(圖片來自網絡)

Operation(算子)

Operation(算子,有時也稱“操作”)用來封裝對於Tensor 的處理邏輯。 同時也是連接TF的前端和後端之間邏輯處理的基本單元,在實際使用中,用戶可以使用keras等上層封裝API 更方便的表達複雜計算邏輯,但是這些上層模塊的內部,最終也會調用各個算子來完成邏輯的表達。

Graph(計算圖)

用戶在TF 前端調用各API 形成的完整計算邏輯,在內部會以dataflow graph 的形式來表達。 在這一有向無環圖(DAG)上,以算子等作為節點,以Tesnor 等作為邊來指明數據的流動路徑。 在graph 上,有些節點是TF 框架自身根據需要添加的,比如,用戶在training算法階段時,只需要調用各種優化器(Optimizer)的minimize方法,TF 自身就會自動找到前向圖中各算子所對應的梯度算子,並按照數學上的鍊式求導法則,構建出反向梯度子圖。

TF 執行過程示例TensorFlow 數據流計算圖(圖片來自TensorFlow 社區)

Session(會話)

Session 主要是在實際執行graph 時對一次執行的上下文進行維護處理。 當用戶調用其run方法時,TF 就會分析為了獲取這一次的計算目標所需要運行的子圖,並結合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 
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);
  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 整體工程架構

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 算子庫中的TfToRttRttToTf兩個用於類型轉換的算子來完成的。

Rosetta在import時的靜態替換

調用Session.run 接口時

我們同樣hook 了Session.run入口,在其內部完成從上一步驟中RttOp算子到SecureOp算子的轉換。 如果用戶使用TensorBoard 工具查看此時的運行圖,就會看到我們在圖上添加了一個和原生TF 計算圖基本同構的新子圖,這個子圖就是由SecureOp構成。

TensorBoard可以查看得到的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 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)等人。 “ Tensorflow:用於大規模機器學習的系統。” 第12屆{USENIX}操作系統設計和實現專題討論會({OSDI} 16)。 2016。

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

系列文章:

隱私AI 工程技術實踐指南:整體介紹