天天看點

Tensorflow代碼解析(一)

摘要

2015年11月9日,Google釋出​​深度學習架構​​​​TensorFlow​​并宣布開源,并迅速得到廣泛關注,在圖形分類、音頻處理、推薦系統和自然語言處理等場景下都被大面積推廣。TensorFlow系統更新快速,官方文檔教程齊全,上手快速且簡單易用,支援Python和C++接口。本文依據對Tensorflow(簡稱TF)白皮書[1]、TF Github[2]和TF官方教程[3]的了解,從系統和代碼實作角度講解TF的内部實作原理。以Tensorflow r0.8.0為基礎,本文由淺入深的闡述Tensor和Flow的概念。先介紹了TensorFlow的核心概念和基本概述,然後剖析了OpKernels子產品、Graph子產品、Session子產品。

1. TF系統架構

1.1 TF依賴視圖

TF的依賴視圖如圖1所示[4],描述了TF的上下遊關系鍊。

Tensorflow代碼解析(一)

圖 1 TensorFlow依賴視圖

TF托管在github平台,有google groups和contributors共同維護。

TF提供了豐富的深度學習相關的API,支援Python和C/C++接口。

TF提供了可視化分析工具Tensorboard,友善分析和調整模型。

TF支援Linux平台,Windows平台,Mac平台,甚至手機移動裝置等各種平台。

1.2 TF系統架構

圖2是TF的系統架構,從底向上分為裝置管理和通信層、資料操作層、圖計算層、API接口層、應用層。其中裝置管理和通信層、資料操作層、圖計算層是TF的核心層。

Tensorflow代碼解析(一)

圖2 TF系統架構

底層裝置通信層負責網絡通信和裝置管理。裝置管理可以實作TF裝置異構的特性,支援CPU、GPU、Mobile等不同裝置。網絡通信依賴gRPC通信協定實作不同裝置間的資料傳輸和更新。

第二層是Tensor的OpKernels實作。這些OpKernels以Tensor為處理對象,依賴網絡通信和裝置記憶體配置設定,實作了各種Tensor操作或計算。Opkernels不僅包含MatMul等計算操作,還包含Queue等非計算操作,這些将在第5章Kernels子產品詳細介紹。

第三層是圖計算層(Graph),包含本地計算流圖和分布式計算流圖的實作。Graph子產品包含Graph的建立、編譯、優化和執行等部分,Graph中每個節點都是OpKernels類型表示。關于圖計算将在第6章Graph子產品詳細介紹。

第四層是API接口層。Tensor C API是對TF功能子產品的接口封裝,便于其他語言平台調用。

第四層以上是應用層。不同程式設計語言在應用層通過API接口層調用TF核心功能實作相關實驗和應用。

1.3 TF代碼目錄組織

圖3是TF的代碼結構視圖,下面将簡單介紹TF的目錄組織結構。

Tensorflow代碼解析(一)

圖3 TF代碼目錄組織結構

Tensorflow/core目錄包含了TF核心子產品代碼。

public: API接口頭檔案目錄,用于外部接口調用的API定義,主要是session.h 和tensor_c_api.h。

client: API接口實作檔案目錄。

platform: OS系統相關接口檔案,如file system, env等。

protobuf: 均為.proto檔案,用于資料傳輸時的結構序列化.

common_runtime: 公共運作庫,包含session, executor, threadpool, rendezvous, memory管理, 裝置配置設定算法等。

distributed_runtime: 分布式執行子產品,如rpc session, rpc master, rpc worker, graph manager。

framework: 包含基礎功能子產品,如log, memory, tensor

graph: 計算流圖相關操作,如construct, partition, optimize, execute等

kernels: 核心Op,如matmul, conv2d, argmax, batch_norm等

lib: 公共基礎庫,如gif、gtl(google模闆庫)、hash、histogram等。

ops: 基本ops運算,ops梯度運算,io相關的ops,控制流和資料流操作

Tensorflow/stream_executor目錄是并行計算架構,由google stream executor團隊開發。

 Tensorflow/contrib目錄是contributor開發目錄。

Tensroflow/python目錄是python API用戶端腳本。

Tensorflow/tensorboard目錄是可視化分析工具,不僅可以模型可視化,還可以監控模型參數變化。

third_party目錄是TF第三方依賴庫。

eigen3: eigen矩陣運算庫,TF基礎ops調用

gpus: 封裝了cuda/cudnn程式設計庫

2. TF核心概念

TF的核心是圍繞Graph展開的,簡而言之,就是Tensor沿着Graph傳遞閉包完成Flow的過程。是以在介紹Graph之前需要講述一下符号程式設計、計算流圖、梯度計算、控制流的概念。

2.1 Tensor

在數學上,Matrix表示二維線性映射,Tensor表示多元線性映射,Tensor是對Matrix的泛化,可以表示1-dim、2-dim、N-dim的高維空間。圖4對比了矩陣乘法(Matrix Product)和張量積(Tensor Contract),可以看出Tensor的泛化能力,其中張量積運算在TF的MatMul和Conv2D運算中都有用到。

Tensorflow代碼解析(一)

圖4 Tensor contract

Tensor在高維空間數學運算比Matrix計算複雜,計算量也非常大,加速張量并行運算是TF優先考慮的問題,如add, contract, slice, reshape, reduce, shuffle等運算。

TF中Tensor的維數描述為階,數值是0階,向量是1階,矩陣是2階,以此類推,可以表示n階高維資料。

TF中Tensor支援的資料類型有很多,如tf.float16, tf.float32, tf.float64, tf.uint8, tf.int8, tf.int16, tf.int32, tf.int64, tf.string, tf.bool, tf.complex64等,所有Tensor運算都使用泛化的資料類型表示。

TF的Tensor定義和運算主要是調用Eigen矩陣計算庫完成的。TF中Tensor的UML定義如圖4。其中TensorBuffer指針指向Eigen::Tensor類型。其中,Eigen::Tensor[5][6]不屬于Eigen官方維護的程式,由貢獻者提供文檔和維護,是以Tensor定義在Eigen unsupported子產品中。

Tensorflow代碼解析(一)

圖5 Tensor資料結構定義

圖5中,Tensor主要包含兩個變量m_data和m_dimension,m_data儲存了Tensor的資料塊,T是泛化的資料類型,m_dimensions儲存了Tensor的次元資訊。

Eigen:Tensor的成員變量很簡單,卻支援非常多的基本運算,再借助Eigen的加速機制實作快速計算,參考章節3.2。Eigen::Tensor主要包含了

一進制運算(Unary),如sqrt、square、exp、abs等。

二進制運算(Binary),如add,sub,mul,div等

選擇運算(Selection),即if / else條件運算

歸納運算(Reduce),如reduce_sum, reduce_mean等

幾何運算(Geometry),如reshape,slice,shuffle,chip,reverse,pad,concatenate,extract_patches,extract_image_patches等

張量積(Contract)和卷積運算(Convolve)是重點運算,後續會詳細講解。

2.2 符号程式設計

程式設計模式通常分為指令式程式設計(imperative style programs)和符号式程式設計(symbolic style programs)。

指令式程式設計容易了解和調試,指令語句基本沒有優化,按原有邏輯執行。符号式程式設計涉及較多的嵌入和優化,不容易了解和調試,但運作速度有同比提升。

這兩種程式設計模式在實際中都有應用,Torch是典型的指令式風格,caffe、theano、mxnet和Tensorflow都使用了符号式程式設計。其中caffe、mxnet采用了兩種程式設計模式混合的方法,而Tensorflow是完全采用了符号式程式設計,Theano和Tensorflow的程式設計模式更相近。

指令式程式設計是常見的程式設計模式,程式設計語言如python/C++都采用指令式程式設計。指令式程式設計明确輸入變量,并根據程式邏輯逐漸運算,這種模式非常在調試程式時進行單步跟蹤,分析中間變量。舉例來說,設A=10, B=10,計算邏輯:

Tensorflow代碼解析(一)

第一步計算得出C=100,第二步計算得出D=101,輸出結果D=101。

符号式程式設計将計算過程抽象為計算圖,計算流圖可以友善的描述計算過程,所有輸入節點、運算節點、輸出節點均符号化處理。計算圖通過建立輸入節點到輸出節點的傳遞閉包,從輸入節點出發,沿着傳遞閉包完成數值計算和資料流動,直到達到輸出節點。這個過程經過計算圖優化,以資料(計算)流方式完成,節省記憶體空間使用,計算速度快,但不适合程式調試,通常不用于程式設計語言中。舉上面的例子,先根據計算邏輯編寫符号式程式并生成計算圖

Tensorflow代碼解析(一)

其中A和B是輸入符号變量,C和D是運算符号變量,compile函數生成計算圖F,如圖6所示。

Tensorflow代碼解析(一)

圖6 符号程式設計的正向計算圖

最後得到A=10, B=10時變量D的值,這裡D可以複用C的記憶體空間,省去了中間變量的空間存儲。

Tensorflow代碼解析(一)

圖 6是TF中的計算流圖,C=F(Relu(Add(MatMul(W, x), b))),其中每個節點都是符号化表示的。通過session建立graph,在調用session.run執行計算。

Tensorflow代碼解析(一)

圖7 TF符号計算圖

和目前的符号語言比起來,TF最大的特點是強化了資料流圖,引入了mutation的概念。這一點是TF和包括Theano在内的符号程式設計架構最大的不同。所謂mutation,就是可以在計算的過程更改一個變量的值,而這個變量在計算的過程中會被帶入到下一輪疊代裡面去。

Mutation是機器學習優化算法幾乎必須要引入的東西(雖然也可以通過immutable replacement來代替,但是會有效率的問題)。 Theano的做法是引入了update statement來處理mutation。TF選擇了純符号計算的路線,并且直接把更新引入了資料流圖中去。從目前的白皮書看還會支援條件和循環。這樣就幾乎讓TF本身成為一門獨立的語言。不過這一點會導緻最後的API設計和使用需要特别小心,把mutation 引入到資料流圖中會帶來一些新的問題,比如如何處理寫與寫之間的依賴。[7]

2.3 梯度計算

梯度計算主要應用在誤差反向傳播和資料更新,是深度學習平台要解決的核心問題。梯度計算涉及每個計算節點,每個自定義的前向計算圖都包含一個隐式的反向計算圖。從資料流向上看,正向計算圖是資料從輸入節點到輸出節點的流向過程,反向計算圖是資料從輸出節點到輸入節點的流向過程。

圖8是2.2節中圖6對應的反向計算圖。圖中,由于C=A*B,則dA=B*dC, dB=A*dC。在反向計算圖中,輸入節點dD,輸出節點dA和dB,計算表達式為dA=B*dC=B*dD, dB=A*dC=A*dD。每一個正向計算節點對應一個隐式梯度計算節點。

Tensorflow代碼解析(一)

圖8 符号程式設計的反向計算圖

反向計算限制了符号程式設計中記憶體空間複用的優勢,因為在正向計算中的計算資料在反向計算中也可能要用到。從這一點上講,粗粒度的計算節點比細粒度的計算節點更有優勢,而TF大部分為細粒度操作,雖然靈活性很強,但細粒度操作涉及到更多的優化方案,在工程實作上開銷較大,不及粗粒度簡單直接。在神經網絡模型中,TF将逐漸側重粗粒度運算。

2.4 控制流

TF的計算圖如同資料流一樣,資料流向表示計算過程,如圖9。資料流圖可以很好的表達計算過程,為了擴充TF的表達能力,TF中引入控制流。

Tensorflow代碼解析(一)

圖9 Graph的資料流

在程式設計語言中,if…else…是最常見的邏輯控制,在TF的資料流中也可以通過這種方式控制資料流向。接口函數如下,pred為判别表達式,fn1和fn2為運算表達式。當pred為true是,執行fn1操作;當pred為false時,執行fn2操作。

Tensorflow代碼解析(一)

TF還可以協調多個資料流,在存在依賴節點的場景下非常有用,例如節點B要讀取模型參數θ更新後的值,而節點A負責更新參數θ,則節點B必須等節點A完成後才能執行,否則讀取的參數θ為更新前的數值,這時需要一個運算控制器。接口函數如下,tf.control_dependencies函數可以控制多個資料流執行完成後才能執行接下來的操作,通常與tf.group函數結合使用。

Tensorflow代碼解析(一)

TF支援的控制算子有Switch、Merge、Enter、Leave和NextIteration等。

TF不僅支援邏輯控制,還支援循環控制。TF使用和MIT Token-Tagged machine相似的表示系統,将循環的每次疊代标記為一個tag,疊代的執行狀态标記為一個frame,但疊代所需的資料準備好的時候,就可以開始計算,進而多個疊代可以同時執行。

作者:​​柒月​​

繼續閱讀