本文作者王思宇，阿里巴巴算法專家，從事深度學(xué)習(xí)算法平臺建設(shè)，TensorFlow 分布式架構(gòu)設(shè)計與大規(guī)模分布式性能優(yōu)化工作，開源 TensorFlow 項目 contributor。

本文轉(zhuǎn)自：互聯(lián)網(wǎng)西門二少 (id: ximen_yushao)

注：建議在閱讀本文時同時梳理代碼~

1. 問題引入

在使用 TensorFlow 構(gòu)建模型時，為了能夠使用 GPU 的 Device，你可能會用到下面的這樣的寫法。

with tf.device('/gpu:0'):
a = tf.get_variable(.....)
b = .......
c = .......

那么，上面代碼中的 a、b 和 c 就真的一定會放在 GPU:0 上嗎？如果 c 不存在 GPU 上的實現(xiàn)會怎么樣？進一步地，有沒有其他約束會讓用戶的設(shè)置失效？

事實上，當(dāng)你打開 session config 的 log_device_placement 選項后，仔細逐個檢查每個 Op 被放置的位置，你會發(fā)現(xiàn)某些 Op 并沒有如你所愿被你控制，而是被“悄悄地”放到別的 Device 上了。

這并不是 Bug，而是 Placer 算法模塊發(fā)揮了保護作用。Placer 算法是 TensorFlow 中 Placement 設(shè)置的最后一道防線。它工作在 TensorFlow 底層，在盡可能滿足用戶訴求的前提下，暗中糾正部分不合理的 Placement。

且聽我從設(shè)計初衷與源碼上，為你娓娓道來。

2. Placement 設(shè)計初衷

受限于單個 Device 的計算能力和存儲大小，模型分片是重要的需求點之一。它的本質(zhì)是將模型和相關(guān)的計算切分到不同的 Device，如此不但可以解決單個 Device 放不下大模型的問題，還有可能帶來計算加速的收益。

在深度學(xué)習(xí)框架方面，顯然在 TensorFlow 上做模型分片比 Caffe 更加容易，這主要得益于 TensorFlow 的 Placement 機制。Placement 是 TensorFlow 引入的特有概念，它指定某個 Op 與具體 Device 的綁定關(guān)系，因此模型分片問題實際上就是該模型上每個 Op 的 Placement 問題。

在 Python 層面，一共存在兩個 API 與 Placement 相關(guān)的接口，它們不但廣泛存在于框架代碼中，還可以被用戶拿來直接使用。

但是用戶指定 Placement 信息存在一定的不可靠性，它與 Op 的實際情況往往存在一定的矛盾，這就需要 TensorFlow 中的 Placer 模塊來解決。

3. Placer 功能描述

Python 構(gòu)完圖之后，請你把 GraphDef 打印出來，我們要關(guān)注每一個 Node 的 NodeDef 結(jié)構(gòu)(如下圖)，這里有兩個地方和 Placement 相關(guān)。

• device 屬性：它顯示指定了這個 Node 應(yīng)該被放在何種 Device 上，它由用戶通過 with tf.device 指定。

• 字符串標(biāo)記 loc:@xxxx：這是 Placement 的約束條件，隱式指明該 Node 的 Placement 應(yīng)該和哪些 Node 保持一致。xxxx 代表某個 Group 的名字，該 Node 應(yīng)該和 Group 名為 xxxx 內(nèi)的所有 Node 的 Placement 保持一致。

可以想象，以上兩個信息可能會出現(xiàn)矛盾的情形。

Placer 不但要處理二者的矛盾，還要通過一些規(guī)則盡可能避免因 Placement 不當(dāng)帶來的性能問題。每個 Node 在經(jīng)過 Placer 處理后都會得到最終的 Placement 信息，它將重新覆蓋 NodeDef 中的 device 屬性內(nèi)容。

所以，通俗地講，Placer 的功能就是推斷并填入所有 NodeDef 的 device 屬性。

4. 一些前驅(qū)內(nèi)容

梳理邏輯時難免會碰到一些為解決這個問題專門設(shè)立的名詞和經(jīng)典的算法，所以建議在閱讀 Placer 模塊相關(guān)內(nèi)容之前先確認已經(jīng)弄清楚下面的東西，避免走一些彎路。

• 顯式 Placement：指用戶通過 with tf.device 直接指定的 Placement 信息，它將寫入上一小節(jié)中 NodeDef 中的 device 屬性。
• 隱式 Placement：指間接指定的 Placement 信息，這個信息與上一小節(jié)中 NodeDef 中的 loc:@xxxx 對應(yīng)。上一節(jié)說到，xxxx 是一個 Group 的名字，該 Group 內(nèi)所有的 Node 都要求具有相同的 Placement 信息，這個 Group 被叫做 Colocation Group，屬于一種約束 (Constraint) 條件。

• Find-Union 算法：并查集算法，Placer 內(nèi)最重要的算法。TensorFlow 通過 Find-Union 算法高效地處理了 Node 的 Colocation 問題。簡單而言，邏輯上，多個具有相同 Colocation Group 的 Node 應(yīng)該被“并”到同一個組中，從而“查”某個 Node 的 Placement 信息時，可以更快速地獲取整組的信息。在實現(xiàn)時，如何設(shè)計更好的數(shù)據(jù)結(jié)構(gòu)，并高效地實施“并”和“查”兩個過程，是并查集算法的核心。

5. Placer決策基本原則

Placer 會根據(jù)會對 Graph 進行一定程度的分析，并結(jié)合用戶的要求對每個 Node 的 Placement 進行微調(diào)，微調(diào)的原則可以概括為下面四點：

• 盡可能滿足用戶要求 (User Requirement First)：每個 Node 的 Placement 會盡量滿足用戶的要求
• 盡可能使用計算更快的設(shè)備 (High Performance Device)：若某個 Node 的 Placement 沒有被用戶指定，則優(yōu)先分配計算更快的設(shè)備
• 保證程序可運行 (Runable)：若某個 Node 不存在用戶要求的 Placement 相關(guān)實現(xiàn)版本，會退而求其次選擇其它實現(xiàn)版本，保障程序可以用

• 盡可能考慮近鄰特性 (Near When Possible)：在做 Placement 的微調(diào)時考慮節(jié)點的近鄰特性，盡可能減少無意義的拷貝

6. 原則原理詳細展開

1. 盡可能滿足用戶要求 (User Requirement First)

用戶要求分為兩種，一種是顯示指定，表現(xiàn)為在 Node 中設(shè)置的 device 信息；另一種是隱式指定，表現(xiàn)為 loc:@xxxx 屬性，即 Colocation Group。

Placer 會根據(jù)用戶這兩方面的要求并結(jié)合實際情況做 Placement 信息補全和微調(diào)。

文章開頭的截圖展示了某個 Node 的 NodeDef 信息，它表明類型為 MatMul 的 Op 被用戶顯示指定放到 '/device:GPU:0' 上，同時希望放入名為 global_step 的 Colocation Group 中。

NodeDef 中的 device 屬性和 loc:@xxxx 屬性分別由下面兩個 Python 級別的 API 引入，它們都由用戶來控制，有些被用在高層 API 內(nèi)部封裝中。

# device attributes
@tf_export("device")
def device(device_name_or_function):

# colocation attributes
@tf_export("colocate_with")
def colocate_with(op, ignore_existing=False):
2. 盡可能使用更快的計算設(shè)備 (High Performance Device)

如果某個 Node 的 device 屬性中不含 device_type(即 GPU 或 CPU)，那么 Placer 必須決定使用何種 Device。每種 Device 注冊到 TensorFlow 中時都帶有優(yōu)先級，通常高優(yōu)先級的 Device 具有更好的計算性能。

當(dāng)某個 Op 具有多種 Device 實現(xiàn)時，Placer 將選取優(yōu)先級最高的 Device 實現(xiàn)版本，通過設(shè)置 device_type 為所有實現(xiàn)版本中最高優(yōu)先級的 Device 來實現(xiàn)這種選取。

3. 保證程序可運行 (Runable)

這是通過 Soft Placement 機制保證的(在 session config 里可以設(shè)置)。

如果某個 Node 被顯示指定精確放在某 Device 上，但系統(tǒng)中卻沒有該 Device 上的實現(xiàn)版本，那么為了保證程序可用，Soft Placement 將發(fā)揮作用，它將忽略 device type，在系統(tǒng)中按照 Device 優(yōu)先級選取另一個可用的實現(xiàn)版本重新改寫 Placement。

舉例而言，假設(shè)某 Node 的 op 是 SparseToDense，device_type 被指定為 GPU，但目前 SparseToDense 在 TensorFlow 中只有 CPU 的實現(xiàn)，那么 Soft Placement 將改寫該 Node 的 device_type 為 CPU。?

4. 盡可能考慮近鄰特性 (Near When Possible)這塊就比較復(fù)雜了，但我們要抓住重點，你就不會亂：關(guān)注三類特殊的 Op 類型，他們的特殊性，決定了其近鄰是需要特殊處理的，分別是：

• Generator 類 Op：入度為 0，出度為 1 的 Op
• MetaData 類 Op：直接在 Tensor 的元數(shù)據(jù) MetaData 上操作，不改變 Tensor 本身的內(nèi)容，比如 Reshape)

• Ref 類或 Resource 類：例如 Variable 這種可能發(fā)生賦值的 Op(或者叫左值)

在 Placer 中使用以下三種啟發(fā)式規(guī)則來分別應(yīng)對上面三種特殊的 Op。

若某個 Node 是 GeneratorNode，將其與 Consumer 與其放在同一個 Device 上可以防止無意義的跨 Device 拷貝。這一步在算法中被稱之為啟發(fā)式規(guī)則 A；

若某個 Node 是 MetaDataNode，將其與 Producer 放在相同的 Device上也可以防止無意義的跨 Device 拷貝。這一步在算法中被稱為啟發(fā)式規(guī)則 B；

若某個 Node 的輸入是 Reference type 或者是 Reource type，那么盡量將其與輸入放在同一個 Colocation Group中(比如 Variable，對其 assign 等操作肯定直接在 Variable 所在之地執(zhí)行即可，如果 Variable 在 A 處，對其的 assign 在 B 處，顯然是不合理的)。算法中沒有為這個步驟起名字，為了方便我們稱之為啟發(fā)式規(guī)則 C。

7. Placer 決策總體流程

總體流程分為四個步驟，下圖展示了宏觀層面的流程圖。其中最后兩個步驟相對較為復(fù)雜，下一節(jié)中將會細化其流程圖。?

8. Placer 分布詳解與關(guān)鍵代碼

注意！本節(jié)看源碼的時候，要注重結(jié)構(gòu)，而不是每個細節(jié)都去糾纏。

第一步 — 根據(jù)外部指定 Colocation 聚合 Group

一般情況下，沒有被用戶指定 Colocation Group 信息的 Node 會被單獨放入一個 Group 中作為唯一的成員，并以該 Node 的 Name 作為 Group 的名字，所以 Graph 中每個 Node 都會有自己的 Colocation Group。

從邏輯上來說，合并多個 Group 是非常簡單的問題，但是這個場景中的 Group 不僅是 Node 的集合，還包含若干屬性，比如某個 Group 的 possible device 表示這個 Group 可用的所有 Device 集合。

因此我們需要一種數(shù)據(jù)結(jié)構(gòu)和算法，幫助我們在合并兩個 Group 時很方便地生成新 Group 及相關(guān)屬性(方便 Union)，并且能夠根據(jù)某個 Node 快速查看所屬 Group 的所有屬性(快速 Find)，這就是 Find-Union 的優(yōu)勢所在。

Find-Union 算法原理將不在這里描述，這里只給出代碼中 Find-Union 用到的基本數(shù)據(jù)結(jié)構(gòu) — Member，它用來描述 Group 的基本信息。在閱讀下段代碼注釋前，需要對 Find-Union 中的樹形結(jié)構(gòu)含義有基本的理解。

// Represents a node in the disjoint node set forest, and the
// accumulated constraints on the device used by that node.
struct Member {
Member() = default;
// The id of the node that is the parent of this one, or its own
// id if it is a root. parent <= 0 indicates that this member is invalid.
int parent = -1;

// A proxy for the depth of the tree that is used to prefer
// connecting smaller trees to larger trees when merging disjoint
// sets.
int rank = 0;

// The intersection of all device types supported by this node,
// and those of all of its children, in priority order
// of the preferred device.
DeviceTypeVector supported_device_types;

// The merged form of the device requested for this node, with
// those of all of its children.
DeviceNameUtils::ParsedName device_name;

// If this node is a root, stores a list of Devices to which this node
// and all of its children have been assigned, or nullptr if this
// has not yet been computed.
std::vector possible_devices;
};

下面的代碼是處理這一步驟的核心代碼。首先創(chuàng)建 ColocationGraph 對象，這是一個處理 Colocation Group 的工具類，里面使用了 Find-Union 算法對 Group 進行聚合。

在調(diào)用 InitiailizeMembers 對 Find-Union 算法的基本數(shù)據(jù)結(jié)構(gòu)進行初始化之后，就直接調(diào)用 ColocationAllNodes 根據(jù)用戶指定的所有 colocation 信息進行聚合。

ColocationGraph colocation_graph(
graph_, devices_,
options_ == nullptr || options_->config.allow_soft_placement(),
default_device_);

TF_RETURN_IF_ERROR(colocation_graph.InitializeMembers());

// 1. First add all of the nodes. Note that steps (1) and (2)
// requires two passes over the nodes because the graph (and hence
// the constraints) may not be acyclic.
TF_RETURN_IF_ERROR(colocation_graph.ColocateAllNodes());
第二步 — 應(yīng)用啟發(fā)式規(guī)則 C(處理 Ref 類 Op Placement)

這一步將對 Colocation Group 進行調(diào)整。在遍歷 Graph 的每個 Node 時，需要根據(jù) Node input 來決定是否將該 Node 所在的 Group 與 Source Node 所在的 Group 合并。

如果 Node 的 input 是 Reference type 或者 DT_RESOURCE(關(guān)于 DT_RESOURCE 一般會在使用 ResourceVariable 時才會碰到。ResourceVariable 與 Variable 相比具有很多新特性，這些特性是 TF2.0 中主推的內(nèi)容。關(guān)于它的優(yōu)勢我們不在這里展開，只對其 Op 的類型做一個說明。

Variable 在 C++ 層面的 Op 類型是 VariableV2，而 ResourceVariable 在 C++ 層面的 Op 類型為 VarHandleOp。后者產(chǎn)生的 Tensor 就是一種 DT_RESOURCE)，那么就嘗試做合并。在合并之前需要做必要的可行性檢查，適當(dāng)?shù)刂鲃訄箦e。比如在合并時除了要考慮這一對節(jié)點的連接以外，還需要考慮這個 Node 的其他輸入是否屬于 Reference type 或者 DT_RESOURCE。這一部分的代碼比較長，但邏輯比較簡單，這里不再展示。

第三步 — 應(yīng)用啟發(fā)式規(guī)則 B(處理 MetaData 類的 Op Placement)

從這一步開始，Placer 才開始真正的為每個 Node 分配 Device，下面的流程圖中展示了這一步驟。

如果當(dāng)前的 Node 的 device 屬性中已經(jīng)有值，那么 Placer 將不再對其做重復(fù)的 assign 操作，直接跳過這個 Node；

如果當(dāng)前 Node 是 GeneratorNode，先將其放入一個名為 second_pass 的 vector 中；

如果不是以上兩種情況，那么該Node正是這一步驟需要處理的對象。先從該 Node 所在的 Colocation Group 中獲取可用的 Devices(獲取會受到 Soft Placement 的影響)作為候選。如果該 node 是 MetaData node，那么會嘗試應(yīng)用啟發(fā)式規(guī)則 B，否則，將分配候選集中優(yōu)先級最高的 Device。

int assigned_device = -1;

// Heuristic B: If the node only operates on metadata, not data,
// then it is desirable to place that metadata node with its
// input.
if (IsMetadata(node)) {
// Make sure that the input device type is in the list of supported
// device types for this node.
const Node* input = (*node->in_edges().begin())->src();
// TODO(vrv): if the input is empty, consider postponing this
// node's assignment to the second pass, so that we handle the
// case where a metadata node's input comes from a backedge
// of a loop.
if (CanAssignToDevice(input->assigned_device_name(), *devices)) {
assigned_device = input->assigned_device_name_index();
}
}

// Provide the default, if necessary.
if (assigned_device == -1) {
assigned_device = graph_->InternDeviceName((*devices)[0]->name());
}

AssignAndLog(assigned_device, node);
第四步 — 應(yīng)用啟發(fā)式規(guī)則 A(處理 Generator 類的 Op Placement)

這一步將對 second_pass 數(shù)組中的所有的 Node 分配 Device，下面的流程圖中展示了這一步驟。

放在 second_pass 中的代碼全部是 GeneratorNode，所以只需要應(yīng)用啟發(fā)式規(guī)則 A 即可，和步驟 3 一樣，啟發(fā)式規(guī)則 A 的應(yīng)用也是嘗試性的，如果實在不能滿足，會直接分配候選 Device 中優(yōu)先級最高的 Device，下面是啟發(fā)式規(guī)則 A 的應(yīng)用部分代碼。

int assigned_device = -1;

// Heuristic A application.
if (IsGeneratorNode(node)) {
const Node* output = (*node->out_edges().begin())->dst();
int output_device_name = output->assigned_device_name_index();

const bool consumers_on_same_device = std::all_of(
node->out_edges().begin(), node->out_edges().end(),
[output_device_name](const Edge* e) {
return e->dst()->assigned_device_name_index() == output_device_name;
});

if (consumers_on_same_device &&
CanAssignToDevice(output->assigned_device_name(), *devices)) {
assigned_device = output_device_name;
}
}

// Provide the default, if necessary.
if (assigned_device == -1) {
assigned_device = graph_->InternDeviceName((*devices)[0]->name());
}

AssignAndLog(assigned_device, node);

至此，所有 Node 的 Placement 信息都已經(jīng)分配并微調(diào)完畢。

9. 總結(jié)

經(jīng)過 Placer 處理的 GraphDef 解決了顯式和隱式 Placement 信息的所有沖突，可謂是最后一道防線。

在 Placer 之后，GraphDef 將被送入 GraphPartitioner 模塊中根據(jù)每個 Node 的 device 做子圖切分，并插入 Send，Recv 以及必要的 ControlFlow 節(jié)點。因此，此步必不可少。

我們也可以看出，Placer 模塊的核心是對 Placement 進行微調(diào)，由于啟發(fā)式規(guī)則相對簡單，性能問題并未完全解決。甚至，我們馬上可以想到，在分布式模式下，粗糙的 Placement 方案會讓作業(yè)性能變得非常差，因為它會引入計算之外的通信開銷。

TensorFlow 高度靈活的 Placement 控制接口，讓模型并行的策略設(shè)計方面具備相當(dāng)大的想象空間，這也是 DL 系統(tǒng)層面研究的熱點之一。而將 Placement 策略自動化，并隱藏到框架中，似乎是用戶十分關(guān)心的問題。這不但可以提高框架的易用性，讓用戶完全專注在模型算法層面，也可以讓初學(xué)者用戶避免寫出性能較差的程序。

但是自動搜索 Placement 最佳策略的難度非常大，因為它要考慮集群通信的帶寬，以及每個 Op 的計算量，是一個與硬件和環(huán)境高度聯(lián)系的復(fù)雜問題。不僅如此，通常深度學(xué)習(xí)模型含有成千上萬個 Node，這使得方案的搜索空間巨大無比。

對于這個問題的解決辦法，目前是百家爭鳴。如果你對策略感興趣，我這里給你推薦一篇 Google 發(fā)表的論文，它利用強化學(xué)習(xí)搜索更好的分片策略。有興趣的同學(xué)可以參考這篇 ICML 的論文：Device Placement Optimization with Reinforcement Learning。

• Device Placement Optimization with Reinforcement Learninghttps://arxiv.org/abs/1706.04972

— 推薦閱讀 —

總結(jié)

以上是生活随笔為你收集整理的tensorflow打印模型结构_社区分享 | 详解 TensorFlow 中 Placement 的最后一道防线 — Placer 算法...的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。