機器學習系統:設計與實現 計算圖

轉自：https://openmlsys.github.io/chapter_computational_graph/index.html

在上一章節中，我們展示了用戶利用機器學習框架所編寫的程序。這些用戶程序包含了對于訓練數據，模型和訓練過程的定義。然而為了運行這些程序，機器學習系統依然需要解決諸多問題，包括：如何高效執行一個復雜的機器學習模型？如何識別出機器學習模型中需要訓練的參數？如何自動計算更新模型所需的梯度？為了解決這些問題，現代機器學習框架實現計算圖*(Computational graph)這一技術。在本章中，我們詳細討論計算圖的基本組成，生成和執行等關鍵設計。本章的學習目標包括：

• 掌握計算圖的基本構成。
• 掌握計算圖靜態生成和動態生成兩種方法。
• 掌握計算圖的常用執行方法。

1 計算圖的設計背景與作用

圖1.1 基于計算圖的架構

早期的機器學習框架主要為了支持基于卷積神經網絡的圖像分類問題。這些神經網絡的拓撲結構簡單（神經網絡層往往通過串行構建），他們的拓撲結構可以用簡單的配置文件來表達（例如Caffe中基于Protocol Buffer格式的模型定義）。隨著機器學習的進一步發展，模型的拓撲日益復雜（包括混合專家，生成對抗網絡，多注意力模型）。這些模型復雜的拓撲結構（例如說，分支結構，帶有條件的if-else循環)會影響模型算子的執行、自動化梯度計算（一般稱為自動微分）以及訓練參數的自動化判斷。為此，我們需要一個更加通用的技術來執行任意機器學習模型。因此，計算圖應運而生。綜合來看，計算圖對于一個機器學習框架提供了以下幾個關鍵作用：

• 對于輸入數據，算子和算子執行順序的統一表達。 機器學習框架用戶可以用多種高層次編程語言（Python，Julia和C++）來編寫訓練程序。這些高層次程序需要統一的表達成框架底層C和C++算子的執行。因此，計算圖的第一個核心作用是可以作為一個統一的數據結構來表達用戶用不同語言編寫的訓練程序。這個數據結構可以準確表述用戶的輸入數據，模型所帶有的多個算子，以及算子之間的執行順序。
• 定義中間狀態和模型狀態。 在一個用戶訓練程序中，用戶會生成中間變量（神經網絡層之間傳遞的激活值和梯度）來完成復雜的訓練過程。而這其中，只有模型參數需要最后持久化，從而為后續的模型推理做準備。通過計算圖，機器學習框架可以準確分析出中間狀態的生命周期（一個中間變量何時生成，以及何時銷毀），從而幫助框架更好的管理內存。
• 自動化計算梯度。 用戶給定的訓練程序僅僅包含了一個機器學習模型如何將用戶輸入（一般為訓練數據）轉化為輸出（一般為損失函數）的過程。而為了訓練這個模型，機器學習框架需要分析任意機器學習模型和其中的算子，找出自動化計算梯度的方法。計算圖的出現讓自動化分析模型定義和自動化計算梯度成為可能。
• 高效程序執行。 用戶給定的模型程序往往是”串行化”地連接起來多個神經網絡層。通過利用計算圖來分析模型中算子的執行關系，機器學習框架可以更好地發現將算子進行異步執行的機會，從而以更快的速度完成模型程序的執行。

2 計算圖的基本構成

計算圖是用來表示深度學習網絡模型在訓練與推理過程中計算邏輯與狀態的工具。計算框架在后端會將前端語言構建的神經網絡模型前向計算與反向梯度計算以計算圖的形式來進行表示。計算圖由基本數據結構張量(Tensor)和基本運算單元算子(Operator)構成。在計算圖中通常使用節點來表示算子，節點間的有向線段來表示張量狀態，同時也描述了計算間的依賴關系。如 圖2.1所示，將 $\mathbf{Z}=relu(\mathbf{X}?\mathbf{Y})$ 轉化為計算圖表示，數據流將根據圖中流向與算子進行前向計算和反向梯度計算來更新圖中張量狀態，以此達到訓練模型的目的。

圖2.1 簡單計算圖

2.1 張量和算子

在計算框架中，基礎組件包含張量和算子，張量是基礎數據結構，算子是基本運算單元。在數學中定義張量是基于向量與矩陣的推廣，涵蓋標量、向量與矩陣的概念。可以將標量理解為零階張量，向量為一階張量，我們熟悉的RGB彩色圖像即為三階張量。在計算框架中張量不僅存儲數據，還存儲數據類型、數據形狀、維度或秩以及梯度傳遞狀態等多個屬性，如表2.1所示，列舉了主要的屬性和功能。

表2.1 張量屬性 張量屬性功能

形狀(shape)	存儲張量的每個維度的長度，如[3,3,3]
維度或秩(dim)	表示張量維度的數量，標量為0，向量為1、矩陣為2
數據類型(dtype)	表示存儲的數 據類型，如bool、int8、int16、float32、float64等
存儲位置(device)	創建張量時可以指定存儲的設備位置，如CPU、GPU等
名字(name)	張量的標識符

張量的形狀是一個重要的屬性，它記錄了每個軸的長度，也就是張量每個維度的元素數量。秩則代表張量的軸數或者階數。張量中通常可以保存布爾類型、浮點數、整型數以及復數和字符串數據。每一個張量都具有唯一的數據類型，在計算過程中會對所有參與運算的張量進行類型檢查，當發現類型不匹配時就會報錯。部分特殊的計算則必須使用指定的數據類型，比如邏輯運算應為布爾類型。在部分計算框架中張量的屬性中包含可以指明張量存儲的設備位置，比如存儲于CPU、GPU等。張量數據的存儲狀態可以分為可變和不可變兩種，不可變張量一般用于用戶初始化的數據或者網絡模型輸入的數據；而可變張量則存儲網絡權重參數，根據梯度信息更新自身數據。

如圖2.2，標量就是一個零階張量，包含單個數值但沒有軸信息。向量即為一階張量，具有一個軸。二階張量具有兩個軸即秩為二。

圖2.2 張量

通常我們使用的張量是”整齊”的，每個軸上的具有相同的元素個數，就像一個”矩形”或者”立方體”。在特定的環境中，也會使用特殊類型的張量，比如不規則張量和稀疏張量，如圖2.3中所示。不規則張量在某個軸上可能具有不同的元素個數，它們支持存儲和處理包含非均勻形狀的數據，在自然語言處理領域，不規則張量可以存儲不同長度文本的信息。稀疏張量則通常應用于圖數據與圖神經網絡中，采用特殊的存儲格式如坐標表格式（Coordinate List， COO），可以高效存儲稀疏數據，節省存儲空間。

圖2.3 張量分類

算子是構成神經網絡的基本計算單元。算子按照功能可以分為張量操作、神經網絡操作、數據流操作和控制流操作等。

• 張量操作 ：包括張量的結構操作和張量的數學運算。張量結構操作有：張量創建、索引切片、維度變換和合并分割等。張量的數學運算包含標量運算、向量運算和矩陣運算。標量運算符的特點是對張量實施逐元素運算。向量運算符只在一個特定軸上運算，將一個向量映射到一個標量或者另外一個向量。矩陣運算包括矩陣乘法、矩陣范數、矩陣行列式、矩陣求特征值、矩陣分解等運算。
• 神經網絡操作 ：包括特征提取、激活函數、損失函數、優化算法等。特征提取是機器學習中的常見操作，核心是提取比原輸入更具代表性的張量，常見的卷積操作就是特征提取算子。激活函數(Activation Function)負責將神經網絡層的輸入映射到輸出端。引入激活函數是為了增加神經網絡模型的非線性，沒有激活函數的每層都相當于矩陣相乘。常見的激活函數包括S型生長曲線(Sigmoid)、線性矯正單元(Rectified Linear Unit, ReLU)等。損失函數(Loss Function)是用來估量模型的預測值與真實值之間的不一致程度。優化算法基于梯度采用不同策略更新參數權值來最小化損失函數，常見的優化算法有隨機梯度下降法(Stochastic Gradient Descent, SGD)、自適應矩估計(Adaptive Moment Estimation, Adam)等。
• 數據流操作 ：包含數據的預處理與數據載入相關算子，數據預處理算子主要是針對圖像數據和文本數據的裁剪填充、歸一化、數據增強等操作。數據載入通常會對數據集進行隨機亂序(Shuffle)、分批次載入(Batch)以及預載入(Prefetch)等操作。數據流操作主要功能是對原始數據進行處理后，轉換為計算框架本身支持的數據格式，并且按照迭代次數輸入給網絡進行訓練或者推理，提升數據載入速度，減少內存占用空間，降低網絡訓練等待時間。
• 控制流操作 ：可以控制計算圖中的數據流向，當表示靈活復雜的模型時需要控制流。使用頻率比較高的控制流算子有條件運算符和循環運算符。控制流操作一般分為兩類，計算框架本身提供的控制流操作符和前端語言控制流操作符。控制流操作不僅會影響神經網絡模型前向運算的數據流向，也會影響反向梯度運算的數據流向。

2.2 計算依賴

在計算圖中，算子之間存在依賴關系，而這種依賴關系影響了算子的執行順序與并行情況。此外在深度學習算法模型中，計算圖是一個有向無環圖，也即在計算圖中造成循環依賴的數據流向是不被允許的。為了理解計算依賴關系并且分析計算圖中循環與循環依賴之間的區別，下面將對計算圖中的計算節點依賴關系進行講解。

圖2.4 計算依賴

如圖2.4中所示，在此簡單的計算圖中，若將 Matmul1 算子移除則該節點無輸出，導致后續的激活函數無法得到輸入，從而計算圖中的數據流動中斷，這表明計算圖中的算子間具有依賴關系并且存在傳遞性。我們對依賴關系進行區分如下：

• 直接依賴 ：節點 ReLU1 直接依賴于節點 Matmul1 ，即如果節點 ReLU1 要執行運算，必須接受直接來自節點 Matmul1 的輸出數據；
• 間接依賴 ：節點 Add 間接依賴于節點 Matmul1 ，即節點 Matmul1 的數據并未直接傳輸給節點 Add ，而是經過了某個或者某些中間節點進行處理后再傳輸給節點 Add ，而這些中間節點可能是節點 Add 的直接依賴節點，也可能是間接依賴節點；
• 相互獨立 ：在計算圖中節點 Matmul1 與節點 Matmul2 之間并無數據輸入輸出依賴關系，所以這兩個節點間相互獨立。

掌握依賴關系后，分析圖2.5可以得出節點 Add 間接依賴于節點 Matmul ，而節點 Matmul 直接依賴于節點 Add ，此時兩個節點互相等待對方計算完成輸出數據，將無法執行計算任務。若我們手動同時給兩個節點賦予輸入，計算將持續不間斷進行，模型訓練將無法停止造成死循環。循環依賴產生正反饋數據流，被傳遞的數值可能在正方向上無限放大，導致數值上溢，或者負方向上放大導致數值下溢，也可能導致數值無限逼近于0，這些情況都會致使模型訓練無法得到預期結果。在構建深度學習模型時，應避免算子間產生循環依賴。

圖2.5 循環依賴

在深度學習計算框架中，表示循環關系通常是以 展開 機制（Unrolling）來實現。當需要實現循環關系時，循環體的計算子圖按照迭代次數進行復制，將代表相鄰迭代輪次的子圖進行串聯，相鄰迭代輪次的計算子圖之間就是直接依賴關系。循環三次的計算圖進行展開如圖2.6。在計算圖中，每一個張量和運算符都具有獨特的標識符，即使是相同的操作運算，在參與不同計算任務時都具有不同的標識符。區分循環關系和循環依賴的關鍵在于，是否兩個獨特標識符之間的運算互相具有直接依賴和相互依賴。循環關系在展開復制計算子圖的時候會給復制的所有張量和運算符賦予新的標識符，區分被復制的原始子圖，以避免形成循環依賴。

圖2.6 循環展開

2.3 控制流

控制流能夠設定特定的順序執行計算任務。若計算圖中無控制流，則每個節點只執行一次，當所有節點按照順序執行完時，計算圖即完成計算。加入控制流后可以讓計算圖中某些節點循環執行任意次數，也可以根據條件判斷選擇某些節點不執行，控制流使得我們可以構建更加靈活和復雜的模型。許多機器學習模型依賴控制流進行訓練和推理，特別是基于遞歸神經網絡和強化學習的模型就依賴于循環遞歸關系和依據數據的條件執行。

為了提高性能、可擴展性和表達能力，計算框架必須支持控制流。目前主流的計算框架中通常使用兩種方式來提供控制流：

• 計算框架控制原語 ：計算框架在內部設計了低級別細粒度的控制原語運算符，通過原語運算符的結合使用來實現控制流，這種實現方式也被稱為圖內方法（In-graph approach）。此類方法的代表就是TensorFlow中的Switch、Merge、Enter、Exit、NextIteration五個原語。TensorFlow通過組合五個原語提供 tf.cond() 和 tf.while_loop() 來實現條件控制和循環控制。
• 前端語言控制流 ：通過高級語言Python、C++的控制流語句來進行計算圖中的控制決策，這類實現方式也被稱為圖外方法（Out-of-graph approach）。計算框架PyTorch、MindSpore中就直接使用Python的控制流，將控制流和數據流之間保持了嚴格的分離。

圖內方法控制流采用框架原語實現，在進行模型編譯、優化與運行時都具備優勢，并且可以準確的判定機器學習模型中計算梯度時需要緩存的變量，提高運行效率，同時由于不依賴外部語言便于部署到不同環境中去。但由于控制原語缺乏進一步的抽象，對于用戶不友好，需要掌握控制原語的使用方法，結合前端語言使用才能描述復雜模型結構。

相對于圖內方法，圖外方法直接使用前端語言控制流則相對更加靈活易用，用戶編寫模型控制時更加便捷直觀，其缺點在于若要將模型進行優化部署，則需要在編譯階段將前端語言的控制流轉化為框架原語描述。

目前在主流的深度學習計算框架中，均提供圖外方法和圖內方法支持。為了便于理解控制流對前向計算與反向計算的影響，后續的講解均使用 圖外方法 實現控制流。常見的控制流包括條件分支與循環兩種。當模型包含控制流操作時，梯度在反向傳播經過控制流時，需要在反向梯度計算圖中也構造生成相應的控制流，才能夠正確計算參與運算的張量梯度。

下面這段代碼描述了簡單的條件控制，我們使用 matmul 表示矩陣乘法算子：

def control(A, B, C, conditional = True):if conditional:y = matmul(A, B)else:y = matmul(A, C)return y

圖2.7 條件控制計算圖

圖2.7描述上述代碼的前向計算圖和反向計算圖。對于具有if-條件的模型，梯度計算需要知道采用了條件的哪個分支，然后將梯度邏輯應用于該分支。在前向計算圖中張量 C 經過條件控制不參與計算，在反向計算時同樣遵守控制流決策，不會計算關于張量 C 的梯度。

當模型中有循環控制時，循環中的操作可以執行零次或者多次。此時采用展開機制，對每一次操作都賦予獨特的運算標識符，以此來區分相同運算操作的多次調用。每一次循環都直接依賴于前一次循環的計算結果，所以在循環控制中需要維護一個張量列表，將循環迭代的中間結果緩存起來，這些中間結果將參與前向計算和梯度計算。下面這段代碼描述了簡單的循環控制，將其展開得到等價代碼后，可以清楚的理解需要維護張量 Xi 和 Wi 的列表。

def recurrent_control(X : Tensor, W : Sequence[Tensor], cur_num = 3):for i in range(cur_num):X = matmul(X, W[i])return X #利用展開機制將上述代碼展開，可得到等價表示 def recurrent_control(X : Tensor, W : Sequence[Tensor]):X1 = matmul(X, W) #為便于表示與后續說明，此處W = W[0], W1 = W[1], W2 = W[2]X2 = matmul(X1, W1)Y = matmul(X2, W2)return Y

如圖2.8描述了上述代碼的前向計算圖和反向計算圖，循環控制的梯度同樣也是一個循環，它與前向循環相迭代次數相同，執行循環體的梯度計算。循環體輸出的梯度值作為下一次梯度計算的初始值，直至循環結束。

圖2.8 循環控制計算圖

2.4 基于鏈式法則計算梯度

在上一小節循環展開的例子中，當神經網絡接收輸入張量 $\mathbf{Y}$ 后，輸入數據根據計算圖逐層進行計算并保存中間結果變量，直至經過多層的計算后最終產生輸出 ${\mathbf{Y}}_3$ ，這個過程我們稱之為 前向傳播 （Forward propagation）。在深度神經網絡模型訓練過程中，前向傳播的輸出結果與標簽值可以產生一個損失函數結果。模型將來自損失函數的數據信息通過計算圖反向流動，執行梯度計算來進行更新訓練參數，這個過程我們稱之為 反向傳播 （Back propagation）。在神經網絡模型中，反向傳播通常使用損失函數關于參數的梯度來進行更新，也可以使用其他信息進行反向傳播，在這里我們僅討論一般情況。

在這里我們簡單回憶一下復合函數的鏈式法則公式。鏈式法則是微積分中的求導法則，用于求解復合函數中的導數。復合函數的導數是構成復合有限個函數在相應點的導數乘積。假設 f 和 g 是關于實數 x 的映射函數，設 $y=g(x)$ 并且 $z=f(y)=f(g(x))$ ，則 z 對 x 的導數即為：
$$\frac{dz}{dx}=\frac{dz}{dy}\frac{dy}{dx}$$
神經網絡的反向傳播是根據反向計算圖的特定運算順序來執行鏈式法則的算法。由于神經網絡的輸入通常為三維張量，輸出為一維向量。因此將上述復合函數關于標量的梯度法則進行推廣和擴展。假設 $\mathbf{X}$ 是 m 維張量， $\mathbf{Y}$ 為 n 維張量， $\mathbf{z}$ 為一維向量， $\mathbf{Y}=g(\mathbf{X})$ 并且 $\mathbf{z}=f(\mathbf{Y})$ ，則 $\mathbf{z}$ 關于 $\mathbf{X}$ 每一個元素的偏導數即為：
$$\frac{?z}{?x_i}=\sum _j\frac{?z}{?y_i}\frac{?y_i}{?x_i}$$
上述公式可等價地表示為：
$${?}_x\mathbf{z}=(\frac{?\mathbf{Y}}{?\mathbf{X}}{)}^{?}{?}_{\mathbf{Y}}\mathbf{z}$$
其中 $?{\mathbf{X}}_{\mathbf{z}}$ 表示 $\mathbf{z}$ 關于 $\mathbf{X}$ 的梯度矩陣。

上一小節中簡單的循環控制模型前向傳播可以表示為 $\mathbf{Y}={\mathbf{W}}_2({\mathbf{W}}_1(\mathbf{W}(\mathbf{X})))$ 。在反向傳播的過程中可以將前向計算等價為 $\mathbf{Y}={\mathbf{W}}_2{\mathbf{X}}_2$ ，首先得到參數 ${\mathbf{W}}_2$ 的梯度表示。再接著根據 ${\mathbf{X}}_2={\mathbf{W}}_1{\mathbf{X}}_1$ 得到 ${\mathbf{W}}_1$ 的梯度表示，按照層級即可推導得出 $\mathbf{W}$ 的梯度表示。

根據鏈式法則，相應位置的導數乘積即可將網絡得到的損失函數梯度信息傳播到每一個權重參數，應用優化器的參數權重更新規則，即可達到神經網絡模型參數訓練迭代的目的。

根據上述公式我們可以得出循環控制的反向梯度計算過程如下，在下面代碼中偽變量的前綴 grad 代表變量梯度變量， transpose 代表矩陣轉置算子。

grad_X2 = matmul(grad_Y, transpose(W2)) grad_W2 = matmul(transpose(X2), grad_Y) grad_X1 = matmul(grad_X2, transpose(W1)) grad_W1 = matmul(transpose(X1), grad_X2) grad_X = matmul(grad_X1, transpose(W)) grad_W = matmul(transpose(X), grad_X1)

結合公式、代碼以及圖2.9我們可以看出，在反向傳播過程中使用到前向傳播的中間變量。因此保存網絡中間層輸出狀態和中間變量，盡管占用了部分內存但能夠復用計算結果，達到了提高反向傳播計算效率的目的。

圖2.9 反向傳播局部計算圖

在深度學習計算框架中，控制流可以進行嵌套，比如多重循環和循環條件控制，計算圖會對復雜控制流進行準確的描述，以便于執行正確的計算調度與執行任務。

3 計算圖的生成

計算框架執行深度學習模型訓練時，會根據模型結構生成計算圖，通過調度計算圖完成模型計算。在計算框架中可以生成靜態圖和動態圖兩種計算圖。靜態圖對應聲明式編程范式，動態圖對應命令式編程范式。靜態生成可以根據前端語言描述的神經網絡拓撲結構以及參數變量等信息構建一份固定的計算圖，因此靜態圖在執行期間可以不依賴前端語言描述，常用于神經網絡模型的部署，比如移動端人臉識別場景中的應用等。動態圖則需要在每一次執行神經網絡模型依據前端語言描述動態生成一份臨時的計算圖，這意味著計算圖的動態生成過程靈活可變，該特性有助于我們在神經網絡結構調整階段提高效率。主流計算框架TensorFlow、MindSpore均支持動態圖和靜態圖模式；PyTorch則可以通過工具將構建的動態圖神經網絡模型轉化為靜態結構，以獲得高效的計算執行效率。了解兩種計算圖生成方式的優缺點及構建執行特點，可以針對待解決的任務需求，選擇合適的生成方式調用執行神經網絡模型。

3.1 靜態生成

靜態圖的生成與執行原理如圖3.1所示，采用先編譯后執行的方式，該模式將計算圖的定義和執行進行分離。在靜態圖模式下使用前端語言定義模型形成完整的程序表達后，并不使用前端語言解釋器進行執行，而是將前端描述的完整模型交給計算框架。框架在執行模型計算之前會首先對神經網絡模型進行分析，獲取網絡層之間的連接拓撲關系以及參數變量設置、損失函數等信息，接著用一種特殊的靜態數據結構來描述拓撲結構及其他神經網絡模型組件，這種特殊的靜態數據結構通常被稱為靜態計算圖。靜態計算圖可以通過優化策略轉換成等價的更加高效的結構。當進行模型訓練或者推理過程時，靜態計算圖接收數據并通過相應硬件調度執行圖中的算子來完成任務。

圖3.1 靜態圖生成與執行

以構建并執行下列偽代碼，來詳細講解靜態圖的生成與執行， matmul 表示矩陣乘法算子， relu 表示線性矯正單元算子。在部分計算框架中如TensorFlow進行前端定義時，需要聲明并編寫包含數據占位符、損失函數、優化函數、網絡編譯、執行環境以及網絡執行器等在內的預定義配置項，此外還需要使用圖內控制流算子編寫控制語句，代碼較為繁瑣并缺乏可讀性。隨著計算框架設計的改進與發展，框架提供的編程接口和模型構建模式呈現出更加統一和友好的趨勢，比如MindSpore提供動靜態統一的前端編程表達。因此為了便于理解靜態生成的過程與原理，此處使用更加簡潔的語言邏輯描述模型。

def model(X, flag):if flag>0:Y = matmul(W1, X)else:Y = matmul(W2, X)Y = Y + bY = relu(Y)return Y

完成前端語言的模型完整構建表達后，執行模型運算時不會直接接收輸入數據進行計算，而是使用計算框架的編譯器對模型進行編譯。由于在進行靜態生成編譯時并不讀取輸入數據，此時需要一種特殊的張量來表示輸入數據輔助構建完整的計算圖，這種特殊張量就被稱之為”數據占位符”。在上述的偽代碼中輸入數據 X 需要使用占位符在靜態圖中表示。構造偽代碼中的條件控制時，由于在靜態圖模式下構建網絡并沒有執行任何計算，對于條件控制在編譯階段并不會進行邏輯運算完成判斷，因此需要將條件控制算子以及所有的分支計算子圖加入計算圖中。在執行階段網絡接受數據流入，調度條件控制算子時進行邏輯判斷，控制數據流入不同的分支計算子圖中進行后續計算。由于控制流和靜態生成的特殊性，在部分計算框架中前端語言Python的控制流不能夠被正確編譯為等價的靜態圖結構，因此需要使用復雜的圖內方法實現控制流。

在后續的章節中我們會繼續深入了解計算框架靜態生成圖結構的過程。靜態生成的過程是采用計算框架編譯器將代碼編譯為中間表示。計算框架編譯器受傳統編譯器方案啟發，設計體系結構包含兩部分編譯器前端和編譯器后端。中間表示承上啟下貫穿前端和后端，是前端源代碼和目標硬件代碼之間的中間數據格式。在計算框架編譯器中中間表示以計算圖形式存在，編譯器會根據前端神經網絡模型自動構建完整的前向計算圖和反向計算圖。

圖3.2 靜態生成

經過編譯后獲取完整的計算圖，能夠根據全局信息完成圖優化策略，進行編譯優化形成與模型完全等價的靜態圖。編譯器前端負責完成計算圖與硬件無關的轉換和優化，比如算子融合將網絡中的兩個或多個細粒度的算子融合為一個粗粒度算子,比如圖3.2中將 add 算子與 relu 合并為一個操作，可節省中間計算結果的存儲、讀取等過程，降低框架底層算子調度的開銷，從而提升執行性能和效率。編譯器后端負責與硬件相關的計算圖優化、代碼指令生成和編譯，優化手段包括硬件算子選擇、內存分配、內存復用等，提高算子執行效率和內存利用效率，降低內存開銷。編譯器后端因此使用靜態圖模型運行往往能夠獲取更好的性能和更少的內存占用。在后續章節中將詳細介紹更多編譯器前端和編譯器后端的優化策略。

優化完成的計算圖通過編譯器后端根據計算硬件來生成適配的執行代碼。在執行階段，調用執行器接受輸入數據，依據計算圖調度算子執行訓練或者推理任務。在訓練任務調度算子執行時，由于在執行階段已經編譯獲取模型整體結構，計算框架可以利用自動并行算法制定合理的模型切分與并行策略，進一步提高計算效率。

使用靜態圖構建模型，編譯構建完整的計算圖后，計算圖可以進行序列化保存，并且再次執行時允許使用序列化模型直接進行訓練或推理，不需要再次編譯前端語言源代碼。得益于編譯器前端、中間表示、編譯器后端多級的計算框架編譯器體系結構，編譯器后端可以將神經網絡模型中間表示轉換為不同硬件代碼。結合計算圖序列化和計算圖可轉換多種部署硬件代碼兩種特性，使得靜態圖模型可以直接部署在不同的硬件上面，提供高效的推理服務。

盡管靜態圖具備強大的執行計算性能與直接部署能力，但是在部分計算框架中靜態圖模式下，使用前端語言編寫神經網絡模型以及定義模型訓練過程代碼較為繁瑣，尤其掌握圖內控制流方法具備一定的學習難度，因此熟練掌握并使用靜態圖模式對于初學者并不友好。其次，靜態生成采用先編譯后執行的方式，編譯階段和執行階段分離，前端語言構建的神經網絡模型經過編譯后，計算圖結構便固定執行階段不再改變，并且經過優化用于執行的計算圖結構與原始代碼有較大的差距，導致代碼中的錯誤難以定位到準確位置，增加了代碼調試難度。此外在神經網絡模型開發迭代環節，不能即時打印中間結果。若在源碼中添加輸出中間結果的代碼，則需要將源碼重新編譯后，再調用執行器才能獲取相關信息，降低了代碼調試效率。而動態圖模式則擁有更加靈活的特性，接下來講解動態生成機制。

3.2 動態生成

動態圖原理如圖3.3所示，采用解析式的執行方式，其核心特點是編譯與執行同時發生。動態圖采用前端語言自身的解釋器對代碼進行解析，利用計算框架本身的算子分發功能，算子會即刻執行并輸出結果。動態圖模式采用用戶友好的命令式編程范式，使用前端語言構建神經網絡模型更加簡潔。

圖3.3 動態圖原理

由于動態圖模式的編程友好性，動態圖被廣大深度學習研究者青睞使用。接下來使用上一小節的偽代碼來講解動態生成和靜態生成的區別。

盡管靜態圖和動態圖在前端語言表達上略有差異，但本質的區別在于靜態生成和動態生成的編譯執行過程不同。使用前端語言構建完成模型表達后，動態生成并不采用計算框架編譯器生成完整的靜態計算圖，而是采用前端語言的解釋器Python API調用計算框架，框架利用自身的算子分發功能，將Python調用的算子在相應的硬件如CPU、GPU、NPU等上進行加速計算，然后再將計算結果返回給前端。該過程并不產生靜態的計算圖，而是按照前端語言描述模型結構，按照計算依賴關系進行調度執行，動態生成臨時的圖拓撲結構。

圖3.4 動態生成

如圖3.4中所示，神經網絡前向計算按照模型聲明定義的順序進行執行。當模型接收輸入數據 $\mathbf{X}$ 后，計算框架開始動態生成圖拓撲結構，添加輸入節點并準備將數據傳輸給后續節點。模型中存在條件控制時，動態圖模式下會即刻得到邏輯判斷結果并確定數據流向，因此在圖中假設判斷結果為真的情況下，圖結構中僅會添加關于張量 ${\mathbf{W}}_1$ 的 matmul 算子節點。按照代碼制定的模型計算順序與算子依賴關系，計算框架會依次添加 add 算子節點和 ReLU 算子節點。計算框架會在添加節點的同時完成算子分發計算并返回計算結果，同時做好準備向后續添加的節點傳輸數據。當模型再次進行前向計算時，動態生成的圖結構則失效，并再次根據輸入和控制條件生成新的圖結構。相比于靜態生成，可以發現動態生成的圖結構并不能完整表示前端語言描述的模型結構，需要即時根據控制條件和數據流向產生圖結構。由于計算框架無法通過動態生成獲取完整的圖結構，因此動態圖模式下難以進行圖結構優化以提高計算效率。

在靜態生成環節，由于已經獲取完整的神經網絡模型定義，因此可以同時構建出完整的前向計算圖和反向計算圖。而在動態生成中，由于邊解析邊執行的特性，反向梯度計算的構建隨著前向計算調用而進行。在執行前向過程中，計算框架根據前向算子的調用信息，記錄對應的反向算子信息以及參與梯度計算的張量信息。前向計算完畢之后，反向算子與張量信息隨之完成記錄，計算框架會根據前向動態圖拓撲結構，將所有反向過程串聯起來形成整體反向計算圖。最終，將反向圖在計算硬件上執行計算得到梯度用于參數更新。

對應于圖3.4中，當調用到關于張量 ${\mathbf{W}}_1$ 的 matmul 算子節點時，框架會執行兩個操作：調用 matmul 算子，計算關于輸入 $\mathbf{X}$ 和 ${\mathbf{W}}_1$ 的乘積結果，同時根據反向計算過程 $\mathbf{G}\mathbf{r}\mathbf{a}\mathbf{d}\_{\mathbf{W}}_1=\mathbf{G}\mathbf{r}\mathbf{a}\mathbf{d}\_\mathbf{Y}?\mathbf{X}$ ，記錄下需要參與反向計算的算子和張量 X 。計算框架依照算子調度順序記錄參與反向計算的算子和張量。當前向計算執行完畢，計算框架根據動態生成的前向計算圖結構拓撲關系，利用記錄的反向計算算子和張量動態生成反向計算圖，最終完成神經網絡模型的梯度計算和參數更新。

盡管動態生成中完整的網絡結構在執行前是未知的，不能使用靜態圖中的圖優化技術來提高計算執行性能。但其即刻算子調用與計算的能力，使得模型代碼在運行的時候，每執行一句立即進行運算并會返回具體的值，方便開發者在模型構建優化過程中進行錯誤分析、結果查看等調試工作，為研究和實驗提供了高效的助力。

此外得益于動態圖模式靈活的執行計算特性，動態生成可以使用前端語言的原生控制流，充分發揮前端語言的編程友好性特性。解決了靜態圖中代碼難調試、代碼編寫繁瑣以及控制流復雜等問題，對于初學者更加友好，提高了算法開發迭代效率和神經網絡模型改進速率。

3.3 動態與靜態生成的比較

靜態生成和動態生成的過程各有利弊。從使用者的角度可以直觀的感受到靜態圖不能實時獲取中間結果、代碼調試困難以及控制流編寫復雜，而動態圖可以實時獲取結果、調試簡單、控制流符合編程習慣。雖然靜態圖的編寫、生成過程復雜，但是相應的執行性能卻超過動態圖，我們用一個簡單的例子來說明在性能和內存占用方面靜態圖的優勢。

def model(X1, X2):Y1 = matmul(X1, W1)Y2 = matmul(X2, W2)Y = Y1 + Y2output = relu(Y)return output

在靜態生成過程中，計算框架獲取完整的計算圖可以分析出計算 Y1 和 Y2 的過程相對獨立，可以將其進行自動并行計算，加快計算效率。而動態生成的過程中，若無手動配置并行策略，計算框架無法獲取圖結構不能分析出算子之間的獨立性，則只能按照代碼順序執行。模型在輸出結果之前執行了 add 和 relu 算子操作，在靜態生成過程中利用計算圖優化策略中的算子融合方法，可以將這兩個算子融合為一個算子執行，這樣減少了中間變量 Y 的存儲與讀取過程，加快了計算效率，減少了內存占用。而動態生成過程則需要按照順序執行 add 和 relu 兩步操作，需要存儲變量 Y 。除此之外，由于靜態生成能夠同時分析重構出前向計算圖和反向計算圖，可以提前確定反向計算中需要保存的前向中間變量信息。而動態生成則在完成前向計算后才能構建出反向計算圖，為了保證反向計算效率需要保存更多的前向計算中間變量信息，相比之下靜態生成的過程更加節省內存占用。

為了方便讀者對比，將靜態圖和動態圖特性總結見表3.1。

表3.1 靜態圖與動態圖的對比 特性靜態圖動態圖

即時獲取中間結果	否	是
代碼調試難易	難	簡單
控制流實現方式	特定的語法	前端語言語法
性能	優化策略多，性能更佳	圖優化受限，性能較差
內存占用	內存占用少	內存占用相對較多
部署能力	可直接部署	不可直接部署

針對兩種模式的特性，結合任務需求選擇合適的模式可以事半功倍，學術科研以及模型開發調試階段，為了快速驗證思想和迭代更新模型結構可以選擇動態圖模式進行構建算法；網絡模型確定，為了加速訓練過程或者為硬件部署模型，可以選擇靜態圖模式。

3.4 動態圖與靜態圖的轉換與融合

動態圖模式下擁有簡潔的接口和編程體驗，具備友好的調試交互機制。代碼按照編寫順序即時執行，符合我們在編寫模型的直觀感受和習慣。可以快速將算法思想轉化為實際代碼。靜態圖模式下可以分離前后端語言，編譯解析前端語言構建的整體網絡結構，并進行優化后以高效后端語言執行，可以直接用于部署。為了兼顧動態圖易用性和靜態圖部署性能兩方面優勢，目前TensorFlow、MindSpore、PyTorch、PaddlePaddle等主流計算框架均具備動態圖轉靜態圖的功能，支持使用動態圖編寫代碼，框架自動轉換為靜態圖網絡結構。

動態圖轉換為靜態圖的實現方式有兩種：

• 基于追蹤轉換 ：以動態圖模式執行并記錄調度的算子，構建和保存為靜態圖模型。
• 基于源碼轉換 ：分析前端代碼來將動態圖代碼自動轉寫為靜態圖代碼，并在底層自動幫用戶使用靜態圖執行器運行。

基于追蹤轉換 的原理相對簡單，當使用動態圖模式構建好網絡后，使用追蹤（Tracing）進行轉換將分為兩個階段。第一個階段計算框架會創建一個新的計算圖，此時以動態圖模式執行代碼，計算框架會自動追蹤數據流的流動以及算子的調度，將所有的操作捕獲并根據調度順序構建靜態圖模型。第二個階段，當執行完一次動態圖后，計算框架已生成靜態圖，當再次調用相同的模型時，計算框架會自動指向靜態圖模型，以高效的性能執行計算。追蹤技術只是記錄第一次執行動態圖時調度的算子，但若是模型中存在依賴于中間結果的條件分支控制流，只能追蹤到根據第一次執行時觸發的分支。此時構建的靜態圖模型并不是完整的，缺失了數據未流向的其他分支。在后續的調用中，因為靜態模型已無法再改變，若計算過程中數據流向缺失分支會導致模型運行錯誤。同樣的，依賴于中間數據結果的循環控制也無法追蹤到全部的迭代狀態。

動態圖基于前端語言自身的解釋器進行模型代碼的解析執行。比如當Python作為前端語言，采取原生Python邊運行邊解釋的特性，配合框架提供的數據處理/算子分發的功能計算，即可實現動態圖的即時執行特性。而且靜態圖則采用計算框架自帶的圖編譯器，對神經網絡模型進行建圖后，再調用圖結構進行計算。動態圖代碼與靜態圖代碼之間存在差異，不能直接使用靜態圖編譯器，因此基于源碼轉換的方法需要將動態圖代碼轉換為靜態圖代碼描述。

基于源碼轉換 的方式則能夠改善基于追蹤轉換的缺陷。如圖3.5中所示，基于源碼轉換的流程經歷兩個階段。第一個階段，對動態圖模式下的代碼掃描進行詞法分析，通過詞法分析器分析源代碼中的所有字符，對代碼進行分割并移除空白符、注釋等，將所有的單詞或字符都轉化成符合規范的語法單元列表。接著進行語法分析即解析器，將得到的語法單元列表轉換成樹形式，并對語法進行檢查避免錯誤。第二階段，動態圖轉靜態圖的核心部分就是對抽象語法樹進行轉寫，計算框架中對每一個需要轉換的語法都預設有轉換器，每一個轉換器對語法樹進行掃描改寫，將動態圖代碼語法映射為靜態圖代碼語法。其中最為重要的前端語言控制流，會在這一階段分析轉換為靜態圖接口進行實現。轉寫完畢之后，將新的語法樹再還原回靜態圖代碼，就可以使用靜態生成執行。使用該方式可以避免基于追蹤轉換中控制流表達缺失的情況。

圖3.5 基于源碼轉換流程

在使用上述功能的過程中，可以將整體模型動態圖代碼全部轉換為靜態圖代碼，提高計算效率并用于硬件部署。同時也可以將整體模型中的部分函數轉化為局部靜態子圖，靜態子圖會被計算框架視為一個完整的算子并嵌入動態圖中。執行整體動態圖時，當計算到對應的函數會自動調用靜態子圖。使用該方式在一定程度上既保留代碼調試改進的靈活性，又提高了計算效率。

@ms_function #mindspore中基于源碼轉換的函數裝飾器，可以將該函數轉換為靜態圖 def add_and_relu(Y, b):Y = Y + bY = relu(Y)return Ydef model(X, flag):if flag>0:Y = matmul(W1, X)else:Y = matmul(W2, X)Y = add_and_relu(Y, b)return Y

代碼中模型整體可以采用動態生成，而@ms_function可以使用基于源碼轉換的技術將模塊 add_and_relu 的轉化為靜態圖結構。與動態生成中代碼執行相同，模型接受輸入按照模型定義的計算順序進行調度執行，并生成臨時圖結構，當執行語句 Y=add_and_relu(Y,b) 時，計算框架會自動調用該模塊靜態生成的圖結構執行計算。模塊 add_and_relu 可以利用靜態圖中的優化技術來提高計算性能，實現動態圖和靜態圖的混合執行。此外，動靜態轉換的技術常用于模型部署階段，動態圖預測部署時除了需要已經訓練完成的參數文件，還須提供最初的模型組網前端代碼，這使得動態圖部署受到局限性，部署硬件中往往難以提供支持前端語言執行環境。因此當使用動態圖模式訓練完成模型參數后，可以將整體網絡結構轉換為靜態圖格式，將神經網絡模型和參數文件進行序列化保存，與前端代碼完全解耦，擴大模型部署的硬件支持范圍。

主流的計算框架TensorFlow、MindSpore等均提供動靜態相互轉換與融合執行的技術，我們將各框架中支持源碼轉換和追蹤轉換技術的接口梳理如表3.2所示。

表3.2 主流框架動態圖轉換靜態圖支持 框架動態圖轉靜態圖

TensorFlow	@tf_function追蹤算子調度構建靜態圖，其中 AutoGraph機制可以自動轉換控制流為靜態表達
MindSpore	conte xt.set_context(mode=context.PYNATIVE_MODE) 動態圖模式 co ntext.set_context(mode=context.GRAPH_MODE) 靜態圖模式 @ms_function支持基于源碼轉換
PyTorch	torch.jit.script()支持基于源 碼轉換，torch.jit.trace()支持基于追蹤轉換
PaddlePaddle	paddle.jit.to_static()支持基于源碼轉換，paddle.jit.TracedLayer.trace()支持基于追蹤轉換

4 計算圖的調度

模型訓練就是計算圖調度圖中算子的執行過程。宏觀來看訓練任務是由設定好的訓練迭代次數來循環執行計算圖，此時我們需要優化迭代訓練計算圖過程中數據流載入和模型訓練（推理）等多個任務之間的調度執行。微觀上單次迭代需要考慮計算圖內部的調度執行問題，根據計算圖、計算依賴關系、計算控制分析算子的任務調度隊列。優化計算圖的調度和執行性能，目的是為了盡可能充分利用計算資源，提高計算效率，縮短模型訓練和推理時間。接下來會詳細介紹計算圖的調度和執行。

4.1 算子調度執行

算子的執行調度包含兩個步驟，第一個，根據拓撲排序算法，將計算圖進行拓撲排序得到線性的算子調度序列；第二步，將序列中的算子分配到執行流進行運算。算子調度執行的目標是根據計算圖中算子依賴關系，確定算子調度序列，盡可能將序列中的算子并行執行，提高計算資源的利用率。

計算圖中依賴邊和算子構成了一張有向無環圖(Directed Acyclic Graph)，計算框架后端需要將包含這種依賴關系的算子準確地發送到計算資源，比如GPU、NPU上執行。因此，就要求算子需要按照一定的順序排列好再發送給GPU/NPU執行。針對有向無環圖，我們通常使用拓撲排序來得到一串線性的序列。

如圖4.1所示，左邊是一張有向無環圖。圖中包含了a,b,c,d,e五個節點和a->d,b->c,c->d,d->e四條邊(a->d表示d依賴于a，稱之為依賴邊)。將圖的依賴邊表達成節點的入度(圖論中通常指有向圖中某點作為圖中邊的終點的次數之和)，可以得到各個節點的入度信息(a:0, b:0, c:1, d:2, e:1)。拓撲排序就是不斷循環將入度為0的節點取出放入隊列中，直至所有有向無環圖中的節點都加入到隊列中，循環結束。例如，第一步將入度為0的a,b節點放入到隊列中，此時有向無環圖中c,d的入度需要減1，得到新的入度信息(c:0, d:1, e:1)。以此類推，將所有的將所有的節點都放入到隊列中并結束排序。

圖4.1 算子調度執行

生成調度序列之后，需要將序列中的算子與數據分發到指定的GPU/NPU上執行運算。根據算子依賴關系和計算設備數量，可以將無相互依賴關系的算子分發到不同的計算設備，同時執行運算，這一過程稱之為并行計算，與之相對應的按照序貫順序在同一設備執行運算被稱之為串行計算。在深度學習中，當數據集和參數量的規模越來越大，我們在分發數據與算子時通信消耗會隨之而增加，計算設備會在數據傳輸的過程中處于閑置狀態，此時采用同步與異步的任務調度機制可以更好的協調通信與訓練任務，提高通信模塊與計算設備的使用率，在后續的小節中將詳細介紹串行與并行、同步與異步的概念。

4.2 串行與并行

根據任務隊列的執行順序，我們可以將計算圖的任務調度隊列分為以下兩種：

• 串行：隊列中的任務必須按照順序進行調度執行直至隊列結束；
• 并行：隊列中的任務可以同時進行調度執行，加快執行效率。

首先我們從微觀上來分析計算圖內部的串行調度。計算圖中大多數算子之間存在直接依賴或者間接依賴關系，具有依賴關系的算子間任務調度則必定存在執行前后的時間順序。如圖4.2，計算圖接受輸入數據進行前向計算得到預測值，計算損失函數進行反向梯度計算，整體代碼流程后序算子的計算有賴于前序算子的輸出。此時算子的執行隊列只能以串行的方式進行調度，保證算子都能正確接受到輸入數據，才能完成計算圖的一次完整執行。

圖4.2 算子的串行

宏觀上來看迭代訓練之間，每一輪迭代中計算圖必須讀取訓練數據，執行完整的前向計算和反向梯度計算，將圖中所有參數值更新完畢后，才能開始下一輪的計算圖迭代計算更新。所以”數據載入-數據處理-模型訓練”的計算圖整體任務調度是以串行方式進行的。

在分析計算圖內部算子依賴關系時，除了直接依賴和間接依賴之外，存在算子間相互獨立的情況。如圖4.3中op1和op2之間相互獨立，此時可以將兩個算子分配到兩個硬件上進行并行計算。對比串行執行，并行計算可以同時利用更多的計算資源來縮短執行時間。

圖4.3 算子的并行

并行包括算子并行、模型并行以及數據并行。算子并行不僅可以在相互獨立的算子間執行，同時也可以將單個算子合理的切分為相互獨立的兩個子操作，進一步提高并行性。模型并行就是將整體計算圖進行合理的切分，分配到不同設備上進行并行計算，縮短單次計算圖迭代訓練時間。數據并行則同時以不同的數據訓練多個相同結構的計算圖，縮短訓練迭代次數，加快訓練效率。這三種并行方式將在后續章節中進行詳細講解。

4.3 數據載入同步與異步機制

一次完整計算圖的訓練執行過程包含：數據載入、數據預處理、網絡訓練三個環節。三個環節之間的任務調度是以串行方式進行，每一個環節都有賴于前一個環節的輸出。但計算圖的訓練是多輪迭代的過程，多輪訓練之間的三個環節可以用同步與異步兩種機制來進行調度執行。

• 同步：順序執行任務，當前任務執行完后會等待后續任務執行情況，任務之間需要等待、協調運行；
• 異步：當前任務完成后，不需要等待后續任務的執行情況，可繼續執行當前任務下一輪迭代。

以同步機制來執行計算圖訓練時，如圖4.4所示，每一輪迭代中，數據讀取后進行數據預處理操作，然后傳輸給計算圖進行訓練。每一個環節執行完當前迭代中的任務后，會一直等待后續環節的處理，直至計算圖完成一次迭代訓練更新參數值后，才會進行下一輪迭代的數據讀取、數據處理以及網絡訓練。當進行數據載入時，數據處理、模型訓練處于等待的狀態，相反模型處于訓練時，數據載入的I/O通道處于空閑，同步機制造成計算資源和通信資源的浪費。

圖4.5 同步機制

以異步機制來執行計算圖訓練時，如圖4.5所示，在迭代訓練中，當數據通道將數據讀取后交給后續的數據與處理環節后，不需要等待計算圖訓練迭代完成，直接讀取下一批次的數據。對比同步機制，異步機制的引入減少了數據載入、數據預處理、網絡訓練三個環節的空閑等待時間，能夠大幅度縮短循環訓練的整體時間，提高任務執行效率。

圖4.5 異步機制

當我們將異步機制與并行計算結合在一起，如圖4.6所示，利用豐富的計算資源可以進一步提高計算圖訓練效率，縮短訓練時間。

圖4.6 異步并行

5 總結

• 為了兼顧編程的靈活性和計算的高效性，設計了基于計算圖的深度學習框架。
• 計算圖的基本數據結構是張量，基本運算單元是算子。
• 計算圖可以表示機器學習模型的計算邏輯和狀態，利用計算圖分析圖結構并進行優化。
• 計算圖是一個有向無環圖，圖中算子間可以存在直接依賴和間接依賴關系，或者相互關系獨立，但不可以出現循環依賴關系。
• 可以利用控制流來改變數據在計算圖中的流向，常用的控制流包括條件控制和循環控制。
• 計算圖的生成可以分為靜態生成和動態生成兩種方式。
• 靜態圖計算效率高，內存使用效率高，但調試性能較差，可以直接用于模型部署。
• 動態圖提供靈活的可編程性和可調試性，可實時得到計算結果，在模型調優與算法改進迭代方面具有優勢。
• 利用計算圖和算子間依賴關系可以進行模型中的算子執行調度問題。
• 根據計算圖可以找到相互獨立的算子進行并發調度，提高計算的并行性。而存在依賴關系的算子則必須依次調度執行。
• 計算圖的訓練任務可以使用同步或者異步機制，異步能夠有效提高硬件使用率，縮短訓練時間。

總結

以上是生活随笔為你收集整理的机器学习系统:设计与实现 计算图的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。