結(jié)合Rosetta介紹如何定制化改造TensorFlow前后端相關(guān)組件

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

引言

我們之前分享了：【技術(shù)分享】隱私 AI 工程技術(shù)實踐指南，從整體上了介紹了基于深度學(xué)習(xí)框架開發(fā)隱私 AI 框架的?程挑戰(zhàn)和可?解決?案。在本文中，我們將進(jìn)?步結(jié)合 Rosetta 介紹如何定制化改造 TensorFlow 前后端相關(guān)組件，以集成 MPC (Multi-Party Computation) 等隱私計算技術(shù)，同時保留對 TensorFlow 接口 API 的復(fù)?，從?實現(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 框架進(jìn)?步擴展引?隱私計算功能，第?步需要?較深?地了解這些 AI 框架，所以?先讓我們簡單回顧?下 TF 的核?概念以及宏觀的內(nèi)部處理過程。

TensorFlow 核心概念

Tensor（張量）

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

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

Operation（算?）

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

Graph（計算圖）

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

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

Session（會話）

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

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

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

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

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

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

通過后端 CustomC++op，可以以更加?效的?式實現(xiàn)??的邏輯，可以在其中進(jì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

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

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

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

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

Rosetta 對 TensorFlow 的深度定制化

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

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

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

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

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

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

細(xì)?的讀者可以看出，上?在介紹 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)建過程?樣，如果?戶的程序含有模型訓(xùn)練過程，調(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 類型算?的?動求導(dǎo)的，所以 Rosetta 中還對 tf.python.ops.gradients_util 等??進(jìn)?了 hook 改造。?如，在下?這?，我們設(shè)定當(dāng) tensor 的基礎(chǔ)類型為 string 時仍可以繼續(xù)進(jìn)?反向傳播：

反向梯度圖的?動?成

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

補充說明

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

轉(zhuǎn)換時，由于此時知道了即將運?的完整?圖的信息，?如 DAG 圖上有多少了算?需要運?，所以可以在這?進(jìn)??些定制化的優(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)部細(xì)節(jié)，我們會在后續(xù)內(nèi)容中加以介紹。

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

小結(jié)

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

作者介紹

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

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

參考文獻(xiàn)

[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

責(zé)任編輯：xj

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

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