本文來自社區投稿與征集，作者Rosetta 技術團隊。本文源自他們在 GDG 活動中的分享。

引言

我們之前分享了：【技術分享】隱私 AI 工程技術實踐指南，從整體上了介紹了基于深度學習框架開發隱私 AI 框架的?程挑戰和可?解決?案。在本文中，我們將進?步結合 Rosetta 介紹如何定制化改造 TensorFlow 前后端相關組件，以集成 MPC (Multi-Party Computation) 等隱私計算技術，同時保留對 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 形成的完整計算邏輯，在內部會以 data?ow graph 的形式來表達。在這?有向?環圖 (DAG) 上，以算?等作為節點，以 Tesnor 等作為邊來指明數據的流動路徑。在 graph 上，有些節點是 TF 框架?身根據需要添加的，?如，?戶在 training 算法階段時，只需要調?各種優化器 (Optimizer) 的 minimize ?法，TF ?身就會?動找到前向圖中各算?所對應的梯度算?，并按照數學上的鏈式求導法則，構建出反向梯度?圖。

TensorFlow 數據流計算圖

Session（會話）

Session 主要是在實際執? graph 時對?次執?的上下?進?維護處理。當?戶調?其 run ?法時，TF 就會分析為了獲取這?次的計算?標所需要運?的?圖，并結合 TF 內置的強?的并?優化、分布式執?等模塊，將所需要執?的邏輯進?步拆分為各個?圖，各?映射到當前的可?設備資源上，最終調度這些設備以并?的?式?效完成計算任務。

TensorFlow 分布式并?執?（圖?來??絡）

TensorFlow 的 codebase 本身還是很復雜的，篇幅所限，難以在此對 TensorFlow 進?深?的介紹，感興趣的讀者可以參考 TensorFlow 公眾號(TensorFlow_official) 上其他優秀?章以進?步學習 TensorFlow。

TensorFlow 自定義算子庫的擴展方法

TF 提供了?較豐富的擴展?法，除了在 Python 層可以基于內置的豐富算?集合，通過模塊的繼承、組裝等?式得到?定義的功能之外，還可以在后端 C++ 層?定義??的算?[2]。在后端基于 CustomC++op 機制進?擴展相?于在前端層進?擴展有?些特別的優勢：

有時候基于現有 TF 原?算?表達上層?定義邏輯很困難，?在后端實現則更靈活?由；

通過后端 CustomC++op，可以以更加?效的?式實現??的邏輯，可以在其中進?更底層的、?向編譯器等的各種優化；

整體上看，基于 TF 的擴展?具，使? customC++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") .SetIsStateful();

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

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

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

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 整體?程架構

從系統架構和代碼上看，改造的??可以分為兩?部分：

后端 C++ 部分的適配定制。主要以?定義算?的 kernel 形式進?適配。?部分接?的輸?輸出參數是以 tf.string 基礎類型的 Tensor，??封裝的是?定義的密?數據。在隱私算? SecureOps 的 kernel 內部會進?步調?統?的密碼協議接?來完成 TF 到隱私計算功能的聯通。

前端 Python 部分的適配定制。這?除了在 Python 前端引?我們?定義的算?庫之外，還需要進?步改造 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 兩個?于類型轉換的算?來完成的。

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 圖的構建后，接下?就是依賴 TensorFlow ?身的圖執?引擎來調度執?各個 SecureOp 的后端實現了，在這個 kernal 中，為了和具體使?的隱私計算技術解耦，我們所調?的是密碼協議接?，?如 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/Rosettaf)，歡迎關注并參與到 Rosetta 社區中來。

參考文獻

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

[2] TensorFlow 對定制化 Op 擴展的?持：https://tensor?ow.google.cn/guide/create_op

責任編輯：xj

原文標題：社區分享 | ?向隱私 AI 的 TensorFlow 深度定制化實踐

文章出處：【微信公眾號：TensorFlow】歡迎添加關注！文章轉載請注明出處。