《智慧計算系統》第五章 程式設計框架原理(上)課程筆記

陌上花已开發表於2024-07-03

《智慧計算系統》第五章 程式設計框架原理(上)課程影片連結:https://www.bilibili.com/video/BV1Ei421i7Rg

本文源自於B站國科大計算所 智慧計算系統課程官方賬號 所公開上傳的影片,在原有影片之上,提取了關鍵幀、將音訊轉成了文字並進行了校正,以便學習使用。在此,也感謝國科大計算所 智慧計算系統課程官方能夠將該課程公開!

若侵權,請聯絡刪除。

智慧計算系統 第五章 程式設計框架原理(上)

大家好,下面我們來介紹智慧計算系統第五章程式設計框架原理

img

這個是第五章的提綱,首先我們會介紹程式設計框架的設計原則,然後我們會圍繞程式設計框架的主要的四大模組兒,分別去介紹程式設計框架中的計算圖構建、計算圖執行、深度學習編譯以及分散式訓練的主要的原理

img

然後最後是本章小結

Img

概述

首先介紹為什麼要了解程式設計框架原理,瞭解程式設計框架的原理對學習智慧計算系統來說是有很大的意義的,它有助於程式設計師編寫與框架底層更為契合、效能更優的程式碼,從而提升演算法的實現效率

此外,在面對新的演算法和硬體挑戰的時候,程式設計師還能夠具備定製化擴充套件培訓框架以支援以提供支援的這個能力

那麼在一個程式設計框架中,它主要包含了四個模組,分別是計算圖構建模組計算圖執行模組深度學習編譯模組分散式訓練模組

其中計算圖構建模組計算圖執行模組必要的,而深度學習編譯模組分散式訓練模組是為了追求更高的效能而需要

那在本章中,我們首先介紹程式設計框架的設計原則,再介紹程式設計框架內部的整體架構;設計原則會抽象的指導程式設計框架的設計,而整體架構則是對這一指導的一個具體的細化實現

Img

Img

設計原則

不同的程式設計框架,它擁有各自的設計哲學

大體上總結可以分為三種,簡潔性、易用性、高效性

簡潔性是指由框架提供一套抽象的機制,那使用者只需要關心宏觀上的操作,而不需要去關心微觀上的具體實現

像我們這個例子裡面,我們要去實現一個把張量 b 和 a 去做加法,得到 c 這樣的一個結果,其中張量 b 原本在 CPU 上,我們需要把張量 b 先搬運到 GPU 上,然後再去做張量 a 和 b 的這個加法的操作

那麼在使用程式設計框架進行對這個功能進行程式設計的時候,我們只需要顯示的去指定一個張量存放的位置和張量移動的這個時機;但是具體的這個張量應該怎麼樣儲存,怎麼樣在裝置間移動,是不需要使用者去操心的,由可以交由程式設計框架去維護

Img

我們再來看易用性;Pytorch,它自誕生就堅持 Python 優先,它提供了符合 Python 語言和庫管理的一個互動方式,而不是一套需要重新學習的內嵌語言

那 Pytorch 的使用者介面設計思路也一直忠於 Python 的語言慣例和開發習慣

Img

高效性是指程式設計框架應進行充分的最佳化,從而儘量提高使用者應用程式的執行效率

比如說像 Tensorflow 這個程式設計框架,他就提供了這個宣告式的靜態圖程式設計方式,這樣使得程式設計框架可以獲得一個完整的計算圖並進行全域性的最佳化

那現有的程式設計框架大部都大部分都支援深度學習的編譯技術,透過這個多層級的表示最佳化來充分的利用使用者硬體的一個計算的能力

同時也基本上都支援多機多卡條件的一個分散式訓練,從而能夠高效支援大規模深度學習的任務

Img

整體架構

程式設計框架的整體結構主要分成四大模組,分別是計算圖構建模組分散式訓練模組深度學習編譯模組以及計算圖執行模組

計算圖構建模組,它的功能是完成從輸入的使用者程式到程式設計框架內部原始計算圖的轉換過程,是程式設計框架的入口模組

分散式訓練模組,它主要是應對更大規模的神經網路,將訓練推理任務從一臺裝置擴充套件到多臺裝置

深度學習編譯模組,是對計算圖分別進行圖層級和圖運算元層級的編譯最佳化,從而提升單裝置上的執行效率

最後是這個計算圖執行模組,它是將最佳化後的計算圖中的張量和操作對映到指定的裝置上去進行具體執行,並給出程式設計框架的輸出結果

Img

那麼我們來看一下這張圖,使用者透過程式設計框架去編寫使用者程式,編寫完了以後的程式,透過程式設計框架的計算圖構建模組,構建出這個使用者程式對應的這個原始的計算圖

那有了這個原始的計算圖以後,透過程式設計框架中提供的這個分散式訓練機制對原始計算圖在多個計算裝置上去進行拆分,得到不同計算裝置上拆分後的計運算元圖

那對於這個拆分以後的計運算元圖,我們在應用深度學習編譯這個模組,這個深度學習編譯模組兒它會針對計算圖分別進行圖層級和運算元層級的編譯最佳化

得到最佳化後的計算圖,針對這個最佳化後的計算圖,我們使用計算圖執行模組,將最佳化以後的計算圖中的張量和操作對映到指定裝置上進行具體執行,並且給出程式設計框架的輸出結果

接下來我們就針對程式設計框架中的 4 個模組的原理去進行簡單的介紹

Img

計算圖構建

首先介紹計算圖構建模組

Img

計算圖它主要是由兩個基本的元素構成,分別是張量以及張量操作,張量就是我們其中稱謂的tensor,張量操作就是operation

計算圖它是一個有向圖,包含了有向的邊,這些有向邊就指明瞭張量的流動方向

下面這張圖它展示了一個神經網路,從原始碼到構建正向和反向計算圖的過程

正向傳播會構建正向計算圖,反向傳播是利用自動求導的原理來去構建反向計算圖

那麼我們這一小節會首先從正向傳播的計算圖構件講起,在介紹反向傳播時計算圖構件的具體原理

Img

正向傳播

在正向傳播中, 輸入張量經過搭建的神經網路層層計算傳遞,最終獲得計算結果

那麼計算圖它的構建形式主要分成兩種,一種是動態圖,一種是靜態圖

動態圖它的意思是說在執行函式的時候,我會按照函式順序逐條語句的去生成節點,然後立即計算返回結果。它的優點就是易除錯,但是它的效能最佳化空間有限

靜態圖是和動態圖相對的,它是在執行計算之前已經構建好了所有圖上的節點,然後在圖執行的時候才計算整個計算圖,並且返回最終的結果,那麼它的優點和取它的優點就是效能好;那麼它的缺點就是不易除錯

Img

動態圖

在PyTorch 1.X的版本里主要使用的就是這個動態圖的機制,動態圖就是說計算圖在函式執行過程中逐漸構建,那它執行的是一種稱為立即模式(Eager)的這樣的一個計算過程,就是每次呼叫語句就會立即執行計算,在 Pytorch 中使用的就是動態圖的機制,就是每次執行一條語句就會重新構建這個計算圖

那我們來看一下左下這張圖的這個程式碼,前四行程式碼,它分別構建了四個張量,那對應的我們就會構建在計算圖中構建相應的四個節點,WhWxx以及這個prev_h 這樣的四個張量節點

接下來的兩行語句分別是進行了兩個矩陣乘法的計算,得到了這個 h2hi2h 這樣的兩個張量,那麼我們對應的也在計算圖中新增了這個兩個矩陣乘的節點,以及 h2hi2h 這樣的兩個張量的節點

再接下來的兩行程式碼,我們去計算 next_h

然後最後一行程式碼是計算loss相應的,每執行完一條,語句都會在計算圖中構建出對應的計算節點、張量節點以及這個對應的邊

當我們的語句編寫完成以後,我們的計算圖也就構建完成

那麼無論這一系列語句它是否在同一個函式里面,我們在下一次構建的時候,在下一次呼叫的時候都會重新構建,即使這個圖的結構可能和這個圖是完全相同的

Img

靜態圖

那我們再來看靜態圖,靜態圖指的是整個網路的結構會在開始計算前就建立完成計算圖,然後在框架執行時它會接收整個計算圖,而不是單一的語句

在 Tensorflow 1.X中使用的就是靜態圖的機制,它使用若干基本控制流的運算元,其實它一共使用了 5 個基本的控制流運算元,分別是 Switch、Merge、Enter、Exit 和 NextIteration 這五個控制流運算元,它透過這五個控制流的運算元的這個不同的組合能夠去實現一些,比如說像跳轉、像迴圈等等這樣一些比較複雜的控制流的場景

在 Pytorch 2.X的版本里面,同樣也支援靜態圖的機制。在 Pytorch 2.0 中採取了圖捕獲TorchDynamo 的這個技術,將使用者的動態圖轉化為靜態圖

Img

反向傳播

那我們再來看反向傳播,在正向傳播的時候,正向計算得到的這個結果和目標結果之間通常會存在一個損失函式的值

然後我們就會使用損失函式對我們的引數去進行求導,然後得到梯度,並且使用這個梯度去更新我們的引數

在這個反向傳播的過程中, PyTorch 就會根據正向傳播的計算圖自動生成對應的反向計算節點,並且把它加入到計算圖中去,構建成反向傳播計算圖

Img

反向傳播過程中一個必不可少的步驟就是要計算導數,那麼自動微分就是一種計算導數的方法,目前常用的幾種求導的方法包括了手動求導法數值求導法符號求導法自動求導法

  • 手動求導法就是指使用鏈式法則求解出梯度公式,然後根據公式去編寫程式碼,代入數值計算得到梯度結果
  • 數值求導法的含義是直接代入數值來近似求解
  • 符號求導法是直接對代數表示式進行求解,最後代入問題數字,它會有表示式膨脹的問題
  • 自動求導法就是使用者只需要描述前向計算的過程,然後由程式設計框架自動的推導反向計算圖,它會先建立表示式,再帶入數值計算,那自動求導法是當前的這個程式設計框架都會支援的求導的方法

Img

手動求導法

好,我們先來看手動求導法,它就是指手動用鏈式法則求解除梯度公式,然後代入數值得到最終的梯度值,這是一種非常直觀的這個求導的方法,那麼它的缺點就是對於大規模的深度學習演算法,手動用鏈式法則進行梯度計算,並且轉換成計算程式是非常困難的

我們需要手動去編寫梯度的求解程式碼,當模型發生變化的時候,演算法也需要隨之修改

Img

數值求導法

那數值求導法就是指利用導數的原始定義去進行求解,就是我們下面給出來的一個公式,我們使用一個非常小的一個數h,然後把它帶入到這個公式中,就能求得某一點處的這個函式的這個導數,那麼數值求導法它的優點就是易操作,並且可以對使用者隱藏求解的過程,那它的缺點就是計算量大,並且它的求解速度比較慢,同時它可能會引起舍入誤差和截斷誤差

Img

符號求導法

符號求導法是利用求導規則來對錶達式進行自動操作,從而獲得導數

我們會把常見的一些求導規則,比如說加法求導規則、減法求導規則、乘法除法這些求導的規則帶入到表示式上,獲得這個求導以後的表示式,然後最後再帶入我們的這個數值,那麼它的問題就是很容易會帶來表示式膨脹的這個問題,最終導致求解的速度變慢

我們以下面這張圖為例來去介紹表示式膨脹的含義啊

那麼我們有一個\(l_n\),這個\(l_n\)是和\(x\)有關的,隨著\(n\)的變化,它\(l_n\)的表示式會變得更復雜

那麼我們對\(l_n\)去計算\(l_n\)相對於\(x\)的導數,這個第3列表示的是使用符號求導法求得的這個表示式的結果,這個第4列表示的是這個使用手動求導法去計算這個導數的這個結果

那我們會發現當\(n\)等於 4 的時候,使用符號求導法計算出來的這個表示式相比於手動求導法求得的這個結果來說已經是非常的複雜了。所以隨著這個\(n\)逐漸增大,這個求導的結果就會出現表示式膨脹的問題

Img

自動求導法

目前程式設計框架中使用的基本上都是自動求導法,那麼根據我們前面的介紹,數值求導法,它強調一開始直接帶入數值去做近似的求解;符號求導法,它強調直接對代數表示式去進行求解,然後才代入問題數值,那麼它們都有各自的一些問題。

自動求導法,它是介於數值求導和符號求導方法之間的一種方法,那它首先將符號求導法應用於最基本的運算元,比如說冪函式、指數函式、對數函式、三角函式等等,然後再代入數值,保留中間的計算結果,最後再應用於整個函式

Img

整個的計算,它分成兩步執行:首先第一步根據原始函式建立計算圖,資料正向傳播,然後我們會計算出所有的中間節點,並且記錄計算圖中的節點依賴關係

第二步我們會反向遍歷計算圖,計算輸出對每一個節點的導數

那對於前項計算中一個資料連線多個輸出資料的情況,在自動求導法裡會把這些輸出資料相對於該資料的導數去進行一個累加

Img

我們這裡根據一個例子來去介紹這個自動求導法,它的這個計算過程,我們有一個想要去計算\(f(x_1, x_2)=(e^{x_1}+x_2)(x_2+1)\)這樣的一個函式,那麼首先第一步我們會根據這個函式的內容去建立計算圖

建立完計算圖以後,我們會讓輸入資料\(x_1\)\(x_2\)去進行一個正向的傳播,計算出所有的中間節點,\(x_3\), \(x_5\), \(x_4\), \(x_6\) 以及\(y\)(下圖右半部分),並且記錄計算圖中的節點依賴關係

那麼我們在做前向計算的時候,會把每一步的中間計算、中間結果都存下來,因為這個在反向計算的時候是要用到的

Img

那第二步我們就要去進行反向計算,要計算輸出對每一個節點的導數,那我們就看這個左邊下邊的這張圖,我們從輸出開始,向輸入的方向逐個節點去計算它們的這個導數。

備註:反向傳播的圖是怎麼來的呢?在下面會有PyTorch的程式碼,個人理解是:對於正向傳播中的每一條,都會有相應的反向傳播的節點;而不是正向傳播中的每一個節點對應反向傳播的節點,這個圖雖然畫的很清楚,但是感覺還是會產生一定誤導的=_=

右邊這張圖我們從下往上看,首先計算第一個節點,也就是 \(y\)相對於\(x_6\) 這個節點的導數,因為 \(y\) 它等於\(x_6\),所以這個 \(y\) 相對\(x_6\)這個節點的導數就為1

然後我們再看\(x_5\),再看這個 \(y\)\(x_5\) 這個節點的導數

根據鏈式法則, \(y\)\(x_5\) 這個節點的導數,它等於 \(y\)\(x_6\) 這個節點的導數,乘以 \(x_6\)\(x_5\)這個節點的導數(\(\frac{\partial y}{\partial x_5}=\overline{x_5}=\frac{\partial y}{\partial x_6}\cdot\frac{\partial x_6}{\partial x_5}\)),那麼 \(y\)\(x_6\)這個節點導數在上一步中已經計算求得了(\(\frac{\partial y}{\partial x_6}=\overline{x_6}\)),那我們只要計算出 \(x_6\)\(x_5\)這個節點的導數,它等於\(x_4\)\(\frac{\partial x_6}{\partial x_5}=x_4\));那我們就把之前向傳播時儲存的\(x_4\)節點的計算值帶入到這個公式中,就求得了這個 \(y\)\(x_5\) 這個節點的導數

小結:我們透過鏈式法則逐步的就計算出輸出對每一個節點的導數。

Img

那這裡需要特殊關注的,就是\(x_2\)這個節點;它的特殊之處就是:它有兩個分支,分別去連向\(x_5\)這個節點以及\(x_4\)這個節點。那麼我們在計算導數的時候,我們就會分開分別計算:先去計算\(x_5\)這個分支,從這個分支上 \(y\)結點對於\(x_2\)的導數,得到這個\(\overline{x_2}^1\),然後再去計算\(x_4\)這條分支上, \(y\)節點對這個\(x_2\)的導數,這個\(\overline{x_2}^2\);然後我們再去把求得的這兩個分支上的導數進行相加\(\overline{x_2}^1+\overline{x_2}^2\)),就得到了最終\(y\)\(x_2\)這個結點的導數。

那這樣的一種自動求導的方法,它實際上是兼具了資料求導法和符號求導法的優點,它很容易操作,並且不會有表示式膨脹的問題。

像神經網路這種模型,它通常輸入是上萬到上百萬維,那輸出的損失函式只有一維,那這樣的模型我們只需要一遍自動求導的過程,就可以求出輸出對各個輸入的導數

Img

各種求導方式對比

我們對四種求導方法去進行了一個比較

  • 從精度上來看,手動求解法、符號求導法和自動求導法都有比較高的精度

  • 從對圖的遍歷次數上來看,自動求導法它只需要對圖遍歷\(N_O+1\),這裡的\(N_O\)指的就是神經網路層的這個輸出個數。(備註:+1加的應該是正向傳播;有幾個輸出,就要反向傳播幾次,所以是\(N_O\)

所以使用自動求導法,它對於輸入維度較大的情況下,它的效能優勢是非常明顯的

Img

PyTorch中的自動求導

在 Python 之中使用的是AutoGrad這個自動微分引擎,使用者只需要一行程式碼,也就是tensor.backward()就可以呼叫它來去自動的計算梯度,並且進行反向傳播

在 PyTorch 中的這個 AutoGrad模組裡面,這個 backward 函式它的實現步驟主要有三步:第一步是正向圖的解析(獲取反向圖節點和梯度節點),第二步是構建反向計算圖的節點(獲取反向圖的),第三步是進行反向梯度的傳播

Img

那我們先來看第一步正向圖的解析;這個裡面列出來了PyTorch框架中的原始碼部分,對於正向圖解析這一部分,這個程式碼在這部分裡面我們會根據輸入配置了rootsgrads這兩個變數,其中roots是反向傳播的節點的集合,同時也是前向傳播的輸出節點的集合,grads它是反向傳播需要的梯度節點集合

Img

第二步是構建反向計算圖的節點;我們透過遍歷正向圖中的節點,構建反向計算圖中的output_edges:對inputs中的每個元素都獲取它對應的張量,檢查這個張量是否是(正向圖中的)葉節點,也就是說它是否有梯度函式grad_fn;它如果不是葉子結點的話,就表明這個結點存在梯度函式grad_fn,那麼我們就建立一條從梯度函式到該節點的邊;這樣遍歷完所有的輸入,我們就構建出了反向計算圖中所有邊的集合output_edges

Img

第三步就是進行反向梯度傳播,此時反向計算圖已經構建完成了,那麼我們就呼叫這個engine.execute(...)來去進行反向的梯度傳播

Img

計算圖執行模組

下面我們來介紹程式設計框架中的計算圖執行模組

Img

那在這一節裡面,我們會講述如何將給定計算圖中的張量以及操作對映到給定裝置上具體執行的整個過程

我們會首先介紹在程式設計框架中裝置管理的方法,然後介紹張量的實現方法,最後講解計算圖中的運算元是如何完成執行的

Img

裝置管理

首先來看裝置管理,裝置是程式設計框架中計算圖執行時的硬體實體,每個裝置都具體負責計運算元圖中的張量存放和運算元計算

那麼常見的裝置包括通用的處理器,也就是CPU,以及我們領域專用的處理器,比如GPU,或者是深度學習處理器DLP

那通用處理器 CPU 的管理方法是比較簡單的,這裡我們就不進行介紹了

在程式設計框架的開發裡面,主要需要新增的是對領域專用的處理器,也就是 DLP 的這個裝置管理的支援,它主要包含包括了三個模組,分別是裝置操作執行流管理以及裝置管理的模組

Img

PyTorch中的裝置被直接按照型別去進行分類,比如說CPU,CUDA, DLP 等等

一個裝置是由一個型別和一個裝置索引或者是序列來去唯一標識的它;前者指定了機器的型別,比如說是 CPU 或者是 GPU 等等,那後者是在有多個特定型別的計算裝置時去標識特定的計算裝置

下圖中parse_type的實現在:pytorch/c10/core/Device.cpp at v2.4.0-rc7 · pytorch/pytorch (github.com)

Img

為了支援裝置管理,Pytorch它在DeviceGuardImplInterface.h裡面定義了抽象的裝置管理類,叫做DeviceGuardImplInterface。這個抽象的管理類,它提供了對裝置管理統一的一個抽象介面,那在這個DeviceGuardImplInterface類裡面提供了對裝置操作執行流管理事件管理的函式介面設計

執行流是裝置上抽象出來的管理計算任務的軟體概念,用於在領域專用處理器上的異構程式設計模型下,完成裝置上任務執行的下發和同步的操作。那具體來說,下發到同一個執行流中的任務具有序列性,下發到不同執行流上的任務能併發執行;因此在程式設計時候,使用者是可以建立多個執行流,並將計算任務分配到不同的執行流中的,從而達到任務併發執行的效果;典型的執行流管理操作包括了執行流建立執行流同步執行流銷燬

事件也是裝置在軟體層面抽象出來的概念,它主要是用來表示裝置上任務執行的狀態和進展,比如記錄事件之間的時間間隔,從而計算裝置執行時間等;事件管理它主要包括計算建立事件建立、時間記錄和事件銷燬等基本操作

Img

那我們來看一下在這個DeviceGuardImplInterface這個類裡面,我們就可以看到它提供了對裝置操作、執行流管理和事件管理的函式介面設計

https://github.com/pytorch/pytorch/blob/v2.4.0-rc7/c10/core/impl/DeviceGuardImplInterface.h#L57

Img

張量實現

邏輯檢視與物理檢視

再來看張量實現,張量是神經網路演算法裡面使用的基本的資料結構,它也是計算圖裡面的核心概念之一,對應了計算圖中不同張量操作之間傳遞流動的資料

張量它有一些基本的屬性,像形狀、佈局、步長、偏移量、資料型別和裝置等等,這些基本的屬性可以統稱為張量資料結構的邏輯檢視;張量資料結構的邏輯檢視是程式設計框架使用者軟體層面上直接控制和表達的一些基本屬性

那對程式設計框架開發者來說,我還需要去維護張量資料結構的另外一種檢視,就是物理檢視。張量資料結構的物理檢視,它主要包括在裝置上的實體地址空間大小指標資料型別等屬性;物理檢視是程式設計框架底層需要維護的基本屬性,它對程式設計框架使用者來說是不可見

那我們下面這張表就是張量資料結構的邏輯檢視和物理檢視的一個基本屬性的對比

Img

張量資料結構(邏輯檢視、物理檢視示例)

邏輯檢視裡主要是透過兩個關鍵變數,也就是偏移量步長,來去最終確定它所對應的張量在物理檢視實體地址空間的定址方法

我們先來看下左圖,當我們使用A[:, 0]來去索引張量A的第一列資料的時候,這個切片以後的張量的邏輯檢視就變成了一個形狀為[2]的張量[1, 3],然後它的步長為2,它的資料型別是int32

備註:上面的這個形狀是我拿PyTorch試出來的,下面的行切片也是,和課程中老師說的不太一樣。但神奇的是,如果在選切片時,使用的是A[: , 0:1],出來的形狀就是[2,1]了...

>>> a = torch.randn(2,2)
>>> a[:, 0].shape, a[1, :].shape
(torch.Size([2]), torch.Size([2]))
>>> a[:, 0:1].shape, a[1:2, :].shape
(torch.Size([2, 1]), torch.Size([1, 2]))

那這樣的一個切片操作,它其實並沒有隱式的去建立一個新的張量並且複製,而只是提供了原本物理檢視下的一個新的邏輯檢視;那它的物理檢視仍然是這個實體地址空間從0x10這個位置開始連續儲存的一塊資料,但是由於它的步長等於2,所以在進行實體地址空間定址的時候,每訪問一個元素,我都要跳躍兩個元素,從而對應了是0x100x18兩個位置的資料

同理,這個下圖右邊的這個圖,當我們使用張量的索引方法 A[1, :]來去索引張量的第二行資料的時候,這個張量的邏輯檢視就變成了一個形狀為[2]的這個張量;它的步長為1,資料型別仍然是int32;這個時候我們就需要在邏輯檢視中額外引入偏移量的屬性,來記錄這個新的邏輯檢視對應的張量在物理檢視上資料實際開始的位置

Img

PyTorch中的張量抽象

我們再來看一下程式設計框架中對張量的這個支援。右邊這張圖顯示了 PyTorch 程式設計框架中實現張量的資料呼叫關係

那麼張量在PyTorch 之中,它存在一個和張量對應的類Tensor,這個 Tensor 它是繼承了它的基類 TensorBase而來;TensorBase類再進一步去呼叫相關的結構體,最終支援張量在不同的裝置型別,像CPU, GPU 和 DLP 上得到支援

Img

在 Python 中是透過張量的抽象類Tensor以及這個儲存抽象類Storage來去分別表示張量資料結構中的邏輯檢視和物理檢視

其中TensorImpl類,它是張量抽象的實現,包含了維度資訊、步長資訊、資料型別、裝置佈局等邏輯視角的張量資訊

StorageImpl類,它是張量的一個儲存實現,包含了記憶體、指標、資料總數等物理視角的張量資訊,它會呼叫結構體allocator去進行張量資料空間的一個分配

那麼我們來看一下底下這張圖,就顯示了在PyTorch程式設計框架中張量實現的流程

首先Tensor類是由TensorBase類繼承而來的,在TensorBase類中包含了唯一的成員變數impl_,這個相當於一個指向TensorImpl的指標,並且它表達了前述張量資料結構中的邏輯檢視

TensorImpl又利用結構體Storage來去表示和張量儲存相關的資訊

一個Storage代表一個張量的底層,支援資料緩衝區(原影片也不太清楚(?TODO)),並且唯一的擁有一個指標storage_impl_Storage(還是Impl?(?TODO))會呼叫結構體allocator去進行張量資料空間的分配

Img

根據後端裝置的不同, allocator的程式碼實現也會有不同

Img

張量記憶體分配

張量的資料結構,包含了邏輯檢視和物理檢視兩種表示;從邏輯檢視到物理檢視的轉換需要完成對張量的記憶體,也就是對張量進行記憶體管理

那根據裝置的型別不同,張量管理的方式也有不同,對於 CPU 來說,主要採用即時分配的這個方式;對GPU來說,主要採用的是記憶體池分配的方式

Img

即時分配

我們先來看 CPU 中採用的即時分配方式,當需要分配張量記憶體的時候,就立刻從系統中申請一塊合適大小的記憶體,這個就是即時分配;核心程式碼部分就是malloc分配空間以及free釋放空間

Img

記憶體池分配

記憶體池分配,它是一種預先分配一塊固定大小記憶體池,然後在需要時從記憶體池中分配記憶體的策略;分配的記憶體它來自事先分配好的記憶體塊,而不是每次都向系統申請新的記憶體

記憶體池分配不僅僅是提供了簡單的記憶體的預分配,它還具備一個自我維護的能力,能夠靈活地處理記憶體塊的拆分和合並等等操作

那在領域專用處理器,比如 GPU 或者是 DLP 上,需要使用處理器的執行時介面來去管理裝置端記憶體,通常就會透過記憶體池的方式去管理裝置端的張量

這種做法相比於呼叫系統 API 直接分配記憶體來講,它不僅減少了系統呼叫開銷,同時還具備兩個重要的優勢,一個是節約了裝置記憶體的使用,並且它減少了裝置記憶體的碎片化問題

Img

分配好記憶體以後,就可以建立張量。建立張量時需要對張量進行初始化,那這個初始化的步驟是首先根據固定記憶體標誌去選擇相應的分配器allocator;這個 allocator可以是即時分配,也可以是記憶體池分配。

在我們下面這個例子裡,使用的是 CPU 記憶體分配,也就是即時分配

然後我們會呼叫這個_empty_generic去進行具體的實現;這個_empty_generic它實現了通用的張量建立邏輯,也就是說我首先去分配儲存空間,然後去建立StorageImpl類,去建立TensorImpl

其他的硬體後端,比如說 GPU 或者 DLP 上去進行張量初始化的時候,我們只需要去選擇相應的分配器,並且設定對應的引數,然後再呼叫_empty_generic就可以實現

Img

運算元執行

我們再來看運算元的執行計算圖的過程,它可以被分解為每個運算元單獨執行的過程

  • 首先透過計算圖生成一個執行序列,這個序列會確定運算元的執行順序來確保正確的資料流和依賴關係
  • 然後我們針對每個運算元去進行運算元實現,包括前端定義後端實現前後端繫結三個步驟
  • 最後進行分派執行,包括查詢適合給定輸入的運算元實現呼叫相應的實現來執行具體的計算任務

Img

執行序列

我們先來看執行序列的確定。計算圖它描述了運算元之間的依賴關係,透過分析計算圖節點之間的依賴關係就可以獲得運算元的執行序列

這個序列一般可以使用拓撲排序的演算法來獲得,它的具體的流程是:首先我要計算每一個節點的入度然後我使用拓撲排序演算法對計算圖的節點進行一個排序,這個通常會採用深度優先搜尋或者是廣度優先搜尋來去遍歷圖中的節點;最後我們要檢查結果,檢查這個執行序列的長度是否是等於圖中的這個節點數

像我們這個例子中,我有一個計算圖的節點,這個裡面包含了 5 個運算元,那麼我們就使用拓撲排序方法,給出了這 5 個運算元的一個執行序列

那麼需要注意的是,對同一個計算圖,我們可能會得到多個不同的結果;這是因為在求拓撲序列的過程中,可能有多個入度為0的節點,我們可以從中任意選擇下一個執行的節點

Img

運算元實現

我們再來看運算元實現,那在獲得運算元的執行序列以後,每個運算元都需要在程式設計框架中完成對應的運算元實現

在我們的程式設計框架裡面,使用者的介面(也就是前端),和具體的實現(也就是後端)一般會採用不同的程式語言,比如說在PyTorch裡面,它的使用者前端一般是使用Python作為這個前端的程式語言,然後用C++來去作為它的後端的程式語言

運算元實現的流程首先是做前端定義,就是指在程式設計框架中配置運算元資訊,這裡邊包含了運算元的輸入輸出以及相關的介面定義,最後再生成前端的介面

第二步後端實現,就是使用C++或者是其他高階的程式語言來去編寫運算元的底層實現程式碼,完成運算元的計算邏輯部分的實現

最後前後端的一個繫結,在這一步驟裡面,我們程式設計框架把前端定義的運算元和後端的具體實現去進行一個繫結

Img

前端定義

PyTorch提供了一種高效的模式用於管理整個運算元實現模組,我們稱為native_function。那在使用native_function這個模式進行運算元實現的時候,需要修改配置檔案native_functions.yaml以新增運算元配置資訊

這個native_functions.yaml這個函式中包含了幾個欄位:

第一個欄位是funcfunc欄位定義了運算元的名稱輸入輸出的引數型別

第二個欄位是variants,這個欄位表示需要自動生成的高階方法

第三個欄位是dispatch,這個欄位表示該運算元所支援的後端型別和對應的實現函式

那下面這段程式碼示例就顯示了這個native_function函式的格式,我們可以看到裡面包含了這個 func欄位、 variants欄位以及dispatch欄位

Img

我們以 PyTorch 中的PReLU這個運算元實現為例,說明使用 native function 模式在 PyTorch 中實現一個以 CPU 為後端執行的運算元的流程

那我們這張程式碼,它展示了 PyTorch 中PReLU運算元的配置檔案,也就是 native 函式,包含了PReLU運算元實現部分:PReLU正向傳播函式實現和PReLU反向傳播函式實現

這個操作,它接受兩個張量作為輸入,其中self表示輸入張量, weight表示該操作的可學習的引數

dispatch欄位,它表示PReLU支援的後端型別和對應的實現函式

Img

PReLU運算元的前端實現程式碼如圖所示,這段程式碼實現了PyTorch中PReLU類的定義,包括PReLU類的相關描述以及正向計算邏輯和可學習引數

然後需要在配置檔案中新增運算元正向傳播函式反向傳播函式對應關係

我們這段示例程式碼就表明了正向傳播函式_prelu_kernel,它對應的反向傳播函式是_prelu_kernel_backward

Img

後端實現

後端實現步驟包含了運算元的表層實現底層實現兩部分

表層實現

表層實現可以看作不同裝置之間的抽象函式介面底層實現可以看作具體到某個裝置上的實際程式碼實現,表層實現和底層實現的程式碼中均需要包含正向傳播函式和反向傳播函式,二者的結構是類似的

首先我們建立一個空的返回物件,然後透過配置TensorIteratorConfig來去建立一個TensorIterator用於迭代張量操作

這個Iterator,它提供了統一的計算抽象,封裝了正向計算的輸入權重以及反向計算的梯度,接下來呼叫prelu_stub函式進行實現,完成PReLU運算元的運算,最後返回計算結果

這裡需要注意的是,這個程式碼中定義的實現只是一個封裝,它沒有完成真正的實現,還需要根據後端的硬體來去編寫對應的這個底層實現

Img

底層實現

底層實現指的就是具體到某個裝置上的實際程式碼實現,比如說在 CPU 上的實現,或者是 GPU 上的實現,或者是 DLP 上的實現

那底層實現中的這個PReLU kernel 函式和表層實現中的 PReLU stub函式會在前後端繫結中去完成對應,這種將表層實現和底層實現解耦的設計使得為多種後端註冊核心函式時可以複用相關的介面

Img

前後端繫結

我們再來看前後端繫結,在前後端繫結的步驟中,我們為每個後端編寫了相應的底層實現步驟,以實現各種不同的硬體和軟體平臺,並且為每一種輸入情況都提供了對應的後端實現的過載版本

在前後端繫結中, dispatch它在分派機制中扮演著排程和控制的角色,確保在不同的後端環境中能選擇正確的實現方法

這個dispatch它會維護一個分派表,分派表的表項記錄著運算元到具體的後端實現的對應關係

分派表可以視為一個二維網格縱軸表示PyTorch支援的運算元橫軸表示支援的分派鍵;這個分派鍵是和後端相應的識別符號;這個表初始是為空的,那當新增一個運算元到後端實現的對應關係的時候,就需要編寫TORCH_LIBRARY_IMPL這個函式去進行註冊

Img

分派執行

在獲得運算元的執行序列並且實現了對應的運算元以後,就需要對運算元進行分派執行;它指的是在執行時根據輸入張量的型別裝置型別查詢並呼叫合適的運算元實現方式

在分派執行過程中, Dispatcher會首先根據輸入張量和其他的資訊計算出對應的分派鍵,然後由該分派鍵找到相應的核心函式

這裡涉及到幾個概念,運算元,指的是 Dispatcher 的排程物件,它代表了具體的計算任務;分派鍵,它是根據輸入張量和其他資訊計算,它可以簡單的理解為與硬體平臺相關聯的識別符號;核心函式,指的就是特定硬體平臺上實現運算元功能的具體程式碼

Img

相關文章