快捷導(dǎo)航

對(duì)Tensorflow中Device實(shí)例的生成和管理詳解

更新時(shí)間：2020年02月04日 19:03:33 作者：jacob_wjj

今天小編就為大家分享一篇對(duì)Tensorflow中Device實(shí)例的生成和管理詳解，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧

1. 關(guān)鍵術(shù)語(yǔ)描述

kernel

在神經(jīng)網(wǎng)絡(luò)模型中，每個(gè)node都定義了自己需要完成的操作，比如要做卷積、矩陣相乘等。

可以將kernel看做是一段能夠跑在具體硬件設(shè)備上的算法程序，所以即使同樣的2D卷積算法，我們有基于gpu的Convolution 2D kernel實(shí)例、基于cpu的Convolution 2D kernel實(shí)例。

device

負(fù)責(zé)運(yùn)行kernel的具體硬件設(shè)備抽象。每個(gè)device實(shí)例，對(duì)應(yīng)系統(tǒng)中一個(gè)具體的處理器硬件，比如gpu:0 device, gpu:1 device, cpu:0 device。一般來(lái)說(shuō)，每個(gè)device實(shí)例同時(shí)包括處理器資源、內(nèi)存資源。device的抽象支持硬件設(shè)備提供的并行處理能力。

2. device是什么

為方便描述，下面我們把在tensorflow里面運(yùn)行的神經(jīng)網(wǎng)絡(luò)模型都統(tǒng)一稱為graph。

我們知道，tensorflow主要針對(duì)的是跨硬件平臺(tái)、分布式、并發(fā)運(yùn)行的場(chǎng)景，參與運(yùn)算的每個(gè)硬件資源，我們都抽象為device實(shí)例，便于管理。

device的主要職責(zé)：

管理處理器資源，為支持device內(nèi)部的并行計(jì)算，進(jìn)一步將其抽象為thread pool或streams：

cpu：使用thread pool來(lái)管理，thread之間可支持不同程度的并行計(jì)算能力

gpu: 針對(duì)nvidia gpu，使用cuda streams來(lái)管理，根據(jù)不同的gpu型號(hào)，可支持不同數(shù)量的stream做并行計(jì)算

管理內(nèi)存資源：為kernel的運(yùn)行，分配和釋放內(nèi)存，進(jìn)一步抽象為Allocator及其各種子類的實(shí)例來(lái)管理。

主機(jī)內(nèi)存：

cpu kernel 計(jì)算時(shí)需要的內(nèi)存。

gpu kernel的輸出結(jié)果如果要放置到主機(jī)內(nèi)存中時(shí)，gpu kernel也需要申請(qǐng)主機(jī)內(nèi)存。

顯存： gpu kernel 計(jì)算時(shí)需要的內(nèi)存。

3. device的種類及應(yīng)用場(chǎng)景

由于device要抽象的設(shè)備種類較多，我們主要描述一下本地運(yùn)行的cpu device、gpu device實(shí)例類型。先用一個(gè)UML圖來(lái)表示一下各種device抽象類的關(guān)系：

可以看到，cpu device實(shí)例使用的類是GPUCompatibleCPUDevice，主要是在ThreadPoolDevice的基礎(chǔ)上，增加了gpu<-> cpu之間內(nèi)存?zhèn)鬏敂?shù)據(jù)的優(yōu)化措施。

gpu device實(shí)例使用的類是 GPUDevice 。

4. device實(shí)例的關(guān)鍵數(shù)據(jù)結(jié)構(gòu)

我們以常用的cpu device，gpu device為例，用下圖描述一下device實(shí)例的關(guān)鍵數(shù)據(jù)結(jié)構(gòu)：

可以看到每個(gè)device實(shí)例內(nèi)部都具備并行處理的能力：

GPUCompatibleCPUDevice實(shí)例
 將 cpu 的計(jì)算資源抽象為thread pool，以支持多thread之間的并發(fā)執(zhí)行；
 將主機(jī)內(nèi)存抽象為 CPUAllocator 實(shí)例來(lái)進(jìn)行管理，為cpu kernel、gpu kernel提供主機(jī)內(nèi)存的申請(qǐng)、釋放功能；

GPUDevice實(shí)例
 將gpu的計(jì)算資源抽象為streams， 由于目前只支持NVIDIA的gpu，所以這里我們可以看作抽象為cuda streams，多個(gè)cuda streams之間的計(jì)算可以并發(fā)處理；
 通過(guò)GPUBFCAllocator實(shí)例來(lái)管理顯存，為gpu kernel提供顯存的申請(qǐng)、釋放功能。

5. device實(shí)例的創(chuàng)建

系統(tǒng)中可用的device實(shí)例，由session發(fā)起創(chuàng)建，歸屬于session實(shí)例。

device的創(chuàng)建，使用Factory 設(shè)計(jì)模式，session會(huì)調(diào)用所有注冊(cè)的device factory，逐一產(chǎn)出符合條件的device實(shí)例。

以DirectSession實(shí)例創(chuàng)建gpu device、cpu device為例，具體流程如下圖所示。

為方便結(jié)合代碼閱讀，已包含主要的類、函數(shù)調(diào)用路徑：

可以看到，最終產(chǎn)出的gpu device、cpu device實(shí)例，都會(huì)保存至DirectSession實(shí)例的 devices_ 表中，由DirectSession實(shí)例進(jìn)行分配和使用。

6. 在graph運(yùn)行階段device的使用

在graph的創(chuàng)建階段，session為每個(gè)node分配一個(gè)具體的device實(shí)例，同時(shí)為每個(gè)node創(chuàng)建一個(gè)具體的kernel實(shí)例，這個(gè)kernel實(shí)例將會(huì)運(yùn)行在分配的device實(shí)例上。(參見Tensorflow 核心流程剖析 2 – 神經(jīng)網(wǎng)絡(luò)模型的創(chuàng)建和分割)

接下來(lái)，在graph的運(yùn)行階段，session會(huì)依次處理graph中的node，調(diào)度node所分配的device實(shí)例，去運(yùn)行node的kernel實(shí)例。

每個(gè)kernel 在運(yùn)行時(shí)，會(huì)向其分配的device，申請(qǐng)需要的計(jì)算資源、內(nèi)存資源等，完成具體的運(yùn)算操作。

上述流程如下圖所示。

為方便結(jié)合代碼閱讀，已包含主要的類、函數(shù)調(diào)用路徑：