快捷導(dǎo)航

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

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

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

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

kernel

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

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

device

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

2. device是什么

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

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

device的主要職責(zé)：

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

cpu：使用thread pool來管理，thread之間可支持不同程度的并行計算能力

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

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

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

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

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

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

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

由于device要抽象的設(shè)備種類較多，我們主要描述一下本地運(yùn)行的cpu device、gpu device實(shí)例類型。先用一個UML圖來表示一下各種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)：

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

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

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

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

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

device的創(chuàng)建，使用Factory 設(shè)計模式，session會調(diào)用所有注冊的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í)例，都會保存至DirectSession實(shí)例的 devices_ 表中，由DirectSession實(shí)例進(jìn)行分配和使用。

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

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

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

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

上述流程如下圖所示。

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