TensorFlow中的設備管理——Device的創建與注冊機制


背景

[作者:DeepLearningStack,阿里巴巴算法工程師,開源TensorFlow Contributor]

歡迎大家關注我的公眾號,“互聯網西門二少”,我將繼續輸出我的技術干貨~

作為一款優秀的異構深度學習算法框架,TensorFlow可以在多種設備上運行算法程序,包括CPU,GPU,Google開發的TPU等。因為TensorFlow的架構特性非常好,可擴展性很強,所以也支持用戶自定義補充其他計算設備,比如可以接入FPGA甚至是自定義芯片等。雖然在Google發布的TensorFlow white paper中並沒有過多的描述設備管理相關的內容,只是從較高層面上闡述了Device以及Job的命名規則,但是其設備管理模塊確實是處於架構中比較核心的地位。本文將從架構層面出發,詳細闡述當前TensorFlow源碼中關於設備管理的設計思想和相關細節,理解這部分內容不但可以加深對TensorFlow源碼的理解,還可以有能力接入一些自定義的設備。本文是TensorFlow設備管理的第一篇文章,為了能讓讀者更好的切入到TensorFlow源碼閱讀過程中,先從較為簡單地Device的創建和注冊機制開始。讀者也可以一邊對照本文一邊對照源碼進行閱讀和梳理,並歡迎大家提出各種相關的意見和建議。

計算設備(Device)定義

Google在2015年發布的第一版TensorFlow white paper中,從功能角度上闡述了Device的相關內容,我們可以總結出關鍵的幾點如下:

1. 在TensorFlow中的Device有着特殊的命名規則,無論在單機還是分布式任務中,都能依靠命名確定唯一的Device,它是Device的唯一標識符;

2. TensorFlow使用注冊機制將實現多種Device的添加管理;

3. 每個Device自己管理Memory的分配和釋放。

Device的命名一般使用/job:{job_name}/task:{job_id}/device:{type}:{device_id}的格式,這是為了更好的支持分布式任務。例如/job:worker/task:17/device:gpu:3就表示該Device是ID為17的worker上的ID為GPU設備。至於分布式中的相關概念會在其他文章中詳細闡述,在這里我們只需要知道Device的命名可以幫助我們在任務中定位到具體的某個唯一Device即可。

Device的注冊機制

TensorFlow有兩處涉及到了設備管理,一處存在於TensorFlow的core中,另一部分存在於XLA中。本文只會闡述TensorFlow core中的內容,關於XLA部分的講解可以參見其他blog。TensorFlow使用工廠來創建各種各樣的Device,並且幾乎為每一種Device都實現了對應的DeviceFactory。初讀代碼時可能會被各種Device類名搞混,下面先從TensorFlow中已經有的Device類出發,給出各種Device的類說明。

Device相關類圖

TensorFlow對不同種類的Device做了多層級的抽象,下面的類圖是從當前TensorFlow源碼中梳理出的比較重要的部分。

上圖中的每個類(class)都可以在TensorFlow的源碼中找到,因為當前TensorFlow的進化過程比較快,代碼結構並不處於一個十分穩定的狀態,所以上述類圖中的類結構關系可能在未來發生一些變化,這一點從注釋中也可以看出一些端倪,但是大的架構不會發生變化,所以梳理類結構也是十分有意義的。下面將對每個類的作用進行簡單地闡述,讀者在理解這些類的作用以及關系后再去閱讀源碼就會非常清晰了。

1. DeviceAttributes:在TensorFlow源碼中並不能直接找到這個類的c++定義,其實它是由protobuf編譯出來的。其含義也很好理解,是對特定Device屬性的封裝,比如Device的type,存儲的限制等等;

2. DeviceBase:定義了Device用到的基本方法,比較重要的是GetAllocator和MakeTensorFromProto,前者返回存儲器的分配器,后者是從Proto中生成Tensor,該方法必須被重寫;

3. Device:這個類比DeviceBase更加具體,新包含了一些用於計算調用的方法,比如Compute函數就會調用某Op的Compute計算;

4. RemoteDevice:這個類會在分布式時使用,在此暫時不進行闡述;

5. SingleThreadedCpuDevice:這是一個僅有單個線程的CPU Device抽象,它和ThreadPoolDevice不同,只被用於in expensive的Op計算,這樣做的好處是避免了一些thread初始化工作;

6. ThreadPoolDevice:這就是CPU Device的實現;

7. RenamedDevice:Device的封裝類,封裝時會再取一個新的Device name;

8. GPUCompatibleCPUDevice:這也是CPU Device的實現,和ThreadPoolDevice不同的是,它更多的是為了和GPU進行交互而存在,從其使用的CudaHostAllocator就可以看出這一點;

9. BaseGPUDevice,GPUDevice:這兩個類都和GPU Device的實現有關,其中GPU Device類只是在繼承BaseGPUDevice的基礎上重寫了Allocator,但沒有理解這樣設計的深層次原因。

DeviceFactory相關類圖

上文提到過,TensorFlow中的Device是通過注冊機制添加到運行的進程中的。注冊機制在開源代碼中是十分常見的設計技巧,它涉及到了一種非常經典的設計模式——工廠模式。在定義每個Device時,通過利用C++事先定義好的宏(Macro)將類對象主動注冊到工廠中,這樣就可以達到在程序啟動完畢時,工廠里已經儲備有各種各樣所需要的內容。在TensorFlow中存在多處使用工廠模式的例子,比如本文闡述的Device管理,以及Session管理等。在其他開源框架中我們也能夠看到這一模式,比如Caffe中的Layer也使用的是工廠模式。

從源碼中可以看到,TensorFlow在啟動時會調用一系列static的函數,這些函數是通過宏(Macro)展開得到的。對於設備管理模塊來說,每種Device都由對應的Factory負責管理,而每種DeviceFactory會在程序啟動時注冊到全局唯一的static device factory表中。下面的類圖展示了各種DeviceFactory之間的繼承關系。

DeviceFactory的注冊機制

在理清了上一小節中Device相關類的繼承關系和說明之后,對於上圖中各種DeviceFactory之間的繼承關系就很好理解了。這里面比較陌生的類是Registrar,這可以看做是一個帶有模板的控制類,它只負責一件事——各種DeviceFactory向全局表的注冊。在代碼層面,注冊函數的調用是通過宏實現的,該宏通過傳入DeviceFactory的類名稱即可觸發Regsitrar的調用邏輯,在每個DeviceFactory的C++實現文件后面都會引用此宏。下圖形象的展示了注冊的過程。

有了上述的DeviceFactory的注冊后,就可以在使用時根據使用的Device類型,從對應的DeviceFactory中“獲取”想要的Device了。

Device的創建

 有了DeviceFactory之后,我們就可以從Factory中拿到各種各樣的Device了。真正從Factory中取出Device的過程是在Session創建時進行的,調用的函數是DeviceFactory中的static函數AddDevices。它會遍歷全局device factories表中全部的DeviceFactory並取出,然后逐個調用每個具體XXDeviceFactory的CreateDevices函數,將創建的Device放進vector數組中。下面給出一個簡化版的時序圖。

上述的時序圖描述的較為簡單,實際上DeviceFactory的static函數調用AddDevices時會先將CPU Device創建出來,如果沒有可用的CPU Device,那么程序就會直接報錯退出(一般情況下不會發生此類情況)。這是因為TensorFlow需要保證當沒有其他Device存在時,至少還有CPU可以完成整體程序的計算和調度運行。創建CPU Device之后,就會去遍歷所有DeviceFactory,把所有能夠創建的Device全部創建出來放入vector數組中。

總結

本文主要闡述了TensorFlow設備管理模塊中的設備創建於注冊機制,它是TensorFlow進行設備管理的第一步,也是最簡單的部分。想要深入TensorFlow源碼的新人可以先從此模塊開始閱讀,進而熟悉TensorFlow的Coding style。Device的創建和注冊過程觸發於程序運行的初始階段,因為創建Device時使用了工廠模式,所以此處涉及到了各種DeviceFactory的定義和注冊。在TensorFlow的C++代碼中,各種DeviceFactory在實現文件中通過宏主動將自己注冊到全局表中,這樣做的目的不但減少了大量重復的注冊代碼,還與Device的創建解耦合,是一個非常經典常見的編碼技巧。至於Device的創建是在創建Session時才會觸發,該過程十分簡單。


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM