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

引言

我們之前分享了：【技術(shù)分享】隱私 AI 工程技術(shù)實踐指南，從整體上了介紹了基于深度學習框架開發(fā)隱私 AI 框架的?程挑戰(zhàn)和可?解決?案。在本文中，我們將進?步結(jié)合 Rosetta 介紹如何定制化改造 TensorFlow 前后端相關(guān)組件，以集成 MPC (Multi-Party Computation) 等隱私計算技術(shù)，同時保留對 TensorFlow 接口 API 的復?，從?實現(xiàn)我們上?篇?章中所強調(diào)的“系統(tǒng)易?性”。

?前 Rosetta 主要基于 TensorFlow 1.14 CPU 版本加以開發(fā)（以下簡稱 TensorFlow 為 TF），這是因為 TF 1.x ?前在?業(yè)界中實際應(yīng)?較為?泛，?引?動態(tài)圖等?級功能的 TF 2.0，則由于接?不向后兼容等問題，仍沒有得到?規(guī)模落地。后續(xù)我們也將在 Rosetta 本身功能穩(wěn)定的基礎(chǔ)上考慮?持 TF 2.0。下?就讓我們開始吧。

TensorFlow 快速回顧

想要基于 AI 框架進?步擴展引?隱私計算功能，第?步需要?較深?地了解這些 AI 框架，所以?先讓我們簡單回顧?下 TF 的核?概念以及宏觀的內(nèi)部處理過程。

TensorFlow 核心概念

Tensor（張量）

深度學習需要完成對?量?維度復雜數(shù)據(jù)的處理，在 TensorFlow 中，? Tensor 來封裝同?類型數(shù)據(jù)的?維數(shù)組。其中，基礎(chǔ)類型除了各種不同精度的整數(shù)、浮點數(shù)外，還?持 tf.string 類型，這給我們提供了進??定義類型改造的可能性。

?個三維 Tensor（圖?來??絡(luò)）

Operation（算?）

Operation（算?，有時也稱“操作”）?來封裝對于 Tensor 的處理邏輯。同時也是連接 TF 的前端和后端之間邏輯處理的基本單元，在實際使?中，?戶可以使? Keras 等上層封裝 API 更?便的表達復雜計算邏輯，但是這些上層模塊的內(nèi)部，最終也會調(diào)?各個算?來完成邏輯的表達。

Graph（計算圖）

?戶在 TF 前端調(diào)?各 API 形成的完整計算邏輯，在內(nèi)部會以 data?ow graph 的形式來表達。在這?有向?環(huán)圖 (DAG) 上，以算?等作為節(jié)點，以 Tesnor 等作為邊來指明數(shù)據(jù)的流動路徑。在 graph 上，有些節(jié)點是 TF 框架?身根據(jù)需要添加的，?如，?戶在 training 算法階段時，只需要調(diào)?各種優(yōu)化器 (Optimizer) 的 minimize ?法，TF ?身就會?動找到前向圖中各算?所對應(yīng)的梯度算?，并按照數(shù)學上的鏈式求導法則，構(gòu)建出反向梯度?圖。

TensorFlow 數(shù)據(jù)流計算圖

Session（會話）

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

TensorFlow 分布式并?執(zhí)?（圖?來??絡(luò)）

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

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

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

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

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

整體上看，基于 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++ 中具體的實現(xiàn)這個 op 所對應(yīng)的內(nèi)部處理邏輯，這就是所謂的后端 “kernel”。TF 提供了?些?便的基類接?，?戶?般只需要定義?個?類，override 實現(xiàn)其中的 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 這樣的宏，將上?所定義的接?和內(nèi)部的實現(xiàn)給綁定起來。這是因為 TF ?持基于不同的輸?、輸出類型和所運?的底層設(shè)備架構(gòu)來定義同?個算?不同的內(nèi)部實現(xiàn)，所以?戶可以定義多種 kernel 實現(xiàn)，告知給系統(tǒng)什么場景下運?具體哪?個 kernel，在實際運?時，TF 就可以根據(jù)不同的設(shè)備、數(shù)據(jù)流上下?調(diào)?不同的 kernel 來實際執(zhí)?此 op。例如：

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

4. 將你的后端算?庫編譯為?個動態(tài)庫 so ?件后，在 Python 層調(diào)?接?引?此模塊，然后就可以如同調(diào)?原?算??樣的?式來調(diào)?這些?定義算?了。例如：

# 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

如果你需要在模型訓練程序中調(diào)?這個?定義算?，你還需要在 Python 層通過 @ops.RegisterGradient("XXXOp") 來注冊這個算?對應(yīng)的梯度算?，通過這種?式，TF 就可以在?動構(gòu)建反向梯度圖時?動的實現(xiàn)對?定義算?梯度的集成。

Rosetta 利用 TF 這?擴展機制引?兩類算?：中間過渡層 RttOps 算?庫和隱私計算 SecureOps 算?庫，前者是為了?持?向?定義數(shù)據(jù)類型的計算圖的構(gòu)建，后者是為了對接后端隱私計算功能，并在執(zhí)?圖時進?動態(tài)綁定。之所以從設(shè)計上區(qū)分這兩類算?，是因為可以進?步解耦圖的構(gòu)建和圖的執(zhí)?，提供更多的靈活性。引?了這兩個基礎(chǔ)的算?庫之后，就可以進?步的進?整體的改造了。

RttOp 算?庫
與后端 MPC 隱私計算完全?關(guān)的輔助中間層，?系列的“浮標”置位算?，?持?定義Tensor類型。其內(nèi)部默認的實現(xiàn)邏輯是和對應(yīng)的 TF 原?算??樣的。

SecureOp 算?庫
完整的前后端算?庫，注冊了對應(yīng)的梯度函數(shù)；在內(nèi)部實現(xiàn)中調(diào)?隱私協(xié)議層的抽象算?接?實現(xiàn)和 TF 的對接。

Rosetta 對 TensorFlow 的深度定制化

如上?篇?章整體介紹的那樣，作為?向?qū)嶋H?業(yè)落地?標的隱私 AI 框架，Rosetta 對于 TF 的改造原則始終是為了提供更加便于 AI 開發(fā)者使?的上層接?，以及兼顧系統(tǒng)后端隱私協(xié)議的可擴展性。

Rosetta 整體?程架構(gòu)

從系統(tǒng)架構(gòu)和代碼上看，改造的??可以分為兩?部分：

后端 C++ 部分的適配定制。主要以?定義算?的 kernel 形式進?適配。?部分接?的輸?輸出參數(shù)是以 tf.string 基礎(chǔ)類型的 Tensor，??封裝的是?定義的密?數(shù)據(jù)。在隱私算? SecureOps 的 kernel 內(nèi)部會進?步調(diào)?統(tǒng)?的密碼協(xié)議接?來完成 TF 到隱私計算功能的聯(lián)通。

前端 Python 部分的適配定制。這?除了在 Python 前端引?我們?定義的算?庫之外，還需要進?步改造 TF 中的?動求導功能等模塊以實現(xiàn)對于新隱私算?的?動構(gòu)建圖、?動求導的?持。

從對程序的動態(tài)處理?度來看，如前?篇?章所說，Rosetta 是經(jīng)過兩個階段的 Pass，來完成到底層多?協(xié)作的 MPC 處理程序的轉(zhuǎn)換。這??部分基于 TF 的前后端改造都是為了完成 StaticPass 階段的轉(zhuǎn)換，即將原? Tensor 轉(zhuǎn)換為?持?定義密?類型的 RttTensor，將原? Operation 轉(zhuǎn)換為?持 tf.string 格式輸?輸出的 RttOp ，并最終在圖開始啟動時進?步的轉(zhuǎn)換為承載實際 MPC 操作的 SecureOp 。

細?的讀者可以看出，上?在介紹 TF 的 customC++op 擴展機制的同時，我們已經(jīng)展示了如何定義 Rosetta 中的單個新算?。接下來，我們介紹?下如何基于這些算?實現(xiàn)計算圖的分階段轉(zhuǎn)換。

計算圖的轉(zhuǎn)換構(gòu)建過程

引入 rosetta 庫時

?戶在前端執(zhí)? import lattciex.rosetta 之后，Rosetta 就會? RttOp 靜態(tài)替換掉原? TF 中對應(yīng)的原? API 算?，且各個原? Tensor 也會被包裝?層到 RttTensor ，其與原? Tensor 的主要區(qū)別是，其數(shù)據(jù)的基礎(chǔ)類型是 tf.string，且對應(yīng)的計算算?是 RttOp。這種基礎(chǔ)類型的轉(zhuǎn)換是基于 RttOp 算?庫中的 TfToRtt 和 RttToTf 兩個?于類型轉(zhuǎn)換的算?來完成的。

Rosetta在import時的靜態(tài)替換

調(diào)用 Session.run 接口時

我們同樣 hook 了 Session.run ??，在其內(nèi)部完成從上?步驟中 RttOp 算?到 SecureOp 算?的轉(zhuǎn)換。如果?戶使? TensorBoard ?具查看此時的運?圖，就會看到我們在圖上添加了?個和原? TF 計算圖基本同構(gòu)的新?圖，這個?圖就是由 SecureOp 構(gòu)成。

TensorBoard 可以查看得到的 SecureOp 計算圖

和上?介紹的原? TF 中的完整圖構(gòu)建過程?樣，如果?戶的程序含有模型訓練過程，調(diào)?了優(yōu)化器 Optimizer 的 minimize ?法，則我們還需要完成對 SecureOp 的反向梯度圖?動?成的?持。

?先，我們需要注冊各個 SecureOp 算?所對應(yīng)的梯度函數(shù)。?如對于隱私矩陣乘法算? SecureMatMul，我們按照底層梯度的計算邏輯，定義其梯度函數(shù)如下：

@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 來統(tǒng)?承載?定義的密?數(shù)據(jù)類型，? TF 本身是不?持對于 tf.string 類型算?的?動求導的，所以 Rosetta 中還對 tf.python.ops.gradients_util 等??進?了 hook 改造。?如，在下?這?，我們設(shè)定當 tensor 的基礎(chǔ)類型為 string 時仍可以繼續(xù)進?反向傳播：

反向梯度圖的?動?成

通過這些精細的定制化改造，最終就可以實現(xiàn)反向梯度?圖的?動?成，可以極?的降低?戶上?隱私計算的開發(fā)難度。

補充說明

并?所有的算?都需要轉(zhuǎn)換為 SecureOp，這是因為如果?個局部?圖中全部的輸?都是本地的常量（公開的寫定到代碼中的數(shù)據(jù)，?需保護），那么就沒有必要將這個?圖轉(zhuǎn)換為多?協(xié)作的隱私計算?式計算，這樣可以減少不必要的計算時間。

轉(zhuǎn)換時，由于此時知道了即將運?的完整?圖的信息，?如 DAG 圖上有多少了算?需要運?，所以可以在這?進??些定制化的優(yōu)化，?如優(yōu)化底層協(xié)議中多?之間的并發(fā)通訊。

在通過上述過程完成在前端層到 SecureOp 圖的構(gòu)建后，接下?就是依賴 TensorFlow ?身的圖執(zhí)?引擎來調(diào)度執(zhí)?各個 SecureOp 的后端實現(xiàn)了，在這個 kernal 中，為了和具體使?的隱私計算技術(shù)解耦，我們所調(diào)?的是密碼協(xié)議接?，?如 SecureMatMul ?最終通過如下代碼?段來調(diào)?內(nèi)部“隱私計算引擎”。這?的內(nèi)部細節(jié)，我們會在后續(xù)內(nèi)容中加以介紹。

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

小結(jié)

在本篇?章中，我們進?步介紹了 Rosetta 是如何深度適配、定制化改造 TensorFlow 的各個組件以引?隱私計算功能的。與其他隱私 AI 開源框架相?，Rosetta 由于需要同時對 TensorFlow 的前端和后端進?擴展，并且完全復?對上層的 API 接?，所以定制化的程度更加深?。這?的改造是偏向于“系統(tǒng)易?性”這??標的，不需要太多涉及 MPC 等隱私計算技術(shù)。

作者介紹

Rosetta 技術(shù)團隊，?群專注于技術(shù)、玩轉(zhuǎn)算法、追求?效的?程師。Rosetta 是?款基于主流深度學習框架 TensorFlow 的隱私 AI 框架，作為矩陣元公司?規(guī)模商業(yè)落地的重要引擎，它承載和結(jié)合了隱私計算、區(qū)塊鏈和 AI 三種典型技術(shù)。

?前 Rosetta 已經(jīng)在 Github 開源(https://github.com/LatticeX-Foundation/Rosettaf)，歡迎關(guān)注并參與到 Rosetta 社區(qū)中來。

參考文獻

[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

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

文章出處：【微信公眾號：TensorFlow】歡迎添加關(guān)注！文章轉(zhuǎn)載請注明出處。