當前位置:
首頁 > 新聞 > 萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

斯坦福大學博士生與 Facebook 人工智慧研究所研究工程師 Edward Z. Yang 是 PyTorch 開源項目的核心開發者之一。他在 5 月 14 日的 PyTorch 紐約聚會上做了一個有關 PyTorch 內部機制的演講,本文是該演講的長文章版本。

萬字綜述,核心開發者全面解讀PyTorch內部機制

打開今日頭條,查看更多圖片

大家好!今天我想談談 PyTorch 的內部機制。

這份演講是為用過 PyTorch並且有心為 PyTorch 做貢獻但卻被 PyTorch 那龐大的 C++ 代碼庫勸退的人提供的。沒必要說謊:PyTorch 代碼庫有時候確實讓人難以招架。本演講的目的是為你提供一份導航圖:為你講解一個「支持自動微分的張量庫」的基本概念結構,並為你提供一些能幫你在代碼庫中尋路的工具和技巧。我預設你之前已經寫過一些 PyTorch,但卻可能還沒有深入理解機器學習軟體庫的編寫方式。

萬字綜述,核心開發者全面解讀PyTorch內部機制

本演講分為兩部分:在第一部分中,我首先會全面介紹張量庫的各種概念。我首先會談談你們知道且喜愛的張量數據類型,並詳細討論這種數據類型究竟能提供什麼,這能讓我們更好地理解其內部真正的實現方式。如果你是一位 PyTorch 高級用戶,你可能已經熟悉其中大部分材料了。我們也會談到「擴展點(extension points)」的三個概念、布局(layout)、設備(device)和數據類型(dtype),這能引導我們思考張量類的擴展的方式。在 PyTorch 紐約聚會的現場演講中,我略過了有關自動梯度(autograd)的幻燈片,但我在這裡會進行一些講解。

第二部分會闡述真正用 PyTorch 寫代碼時所涉及的基本細節。我會告訴你如何在 autograd 代碼中披荊斬棘、什麼代碼是真正重要的以及怎樣造福他人,我還會介紹 PyTorch 為你寫核(kernel)所提供的所有炫酷工具。

概念

張量

張量是 PyTorch 中的核心數據結構。對於張量直觀上所表示的東西,你可能已有很好的理解:張量是一種包含某種標量類型(比如浮點數和整型數等)的 n 維數據結構。我們可以將張量看作是由一些數據構成的,還有一些元數據描述了張量的大小、所包含的元素的類型(dtype)、張量所在的設備(CPU 內存?CUDA 內存?)

萬字綜述,核心開發者全面解讀PyTorch內部機制

另外還有一個你可能沒那麼熟悉的元數據:步幅(stride)。stride 實際上是 PyTorch 最別緻的特徵之一,所以值得稍微多討論它一些。

萬字綜述,核心開發者全面解讀PyTorch內部機制

張量一個數學概念。但要在我們的計算機中表示它,我們必須為它們定義某種物理表示方法。最常用的表示方法是在內存中相鄰地放置張量的每個元素(這也是術語「contiguous(鄰接)」的來源),即將每一行寫出到內存,如上所示。在上面的案例中,我已經指定該張量包含 32 位的整型數,這樣你可以看到每一個整型數都位於一個物理地址中,每個地址與相鄰地址相距 4 位元組。為了記住張量的實際維度,我們必須將規模大小記為額外的元數據。

所以這幅圖與步幅有什麼關係?

萬字綜述,核心開發者全面解讀PyTorch內部機制

假設我想要讀取我的邏輯表示中位置張量 [0,1] 的元素。我該如何將這個邏輯位置轉譯為物理內存中的位置?步幅能讓我們做到這一點:要找到一個張量中任意元素的位置,我將每個索引與該維度下各自的步幅相乘,然後將它們全部加到一起。在上圖中,我用藍色表示第一個維度,用紅色表示第二個維度,以便你了解該步幅計算中的索引和步幅。進行這個求和後,我得到了 2(零索引的);實際上,數字 3 正是位於這個鄰接數組的起點以下 2 個位置。

(後面我還會談到 TensorAccessor,這是一個處理索引計算的便利類(convenience class)。當你使用 TensorAccessor 時,不會再操作原始指針,這些計算過程已經為你隱藏了起來。)

步幅是我們為 PyTorch 用戶講解方法的基本基礎。舉個例子,假設我想取出一個表示以上張量的第二行的張量:

萬字綜述,核心開發者全面解讀PyTorch內部機制

使用高級的索引支持,我只需寫出張量 [1, :] 就能得到這一行。重要的是:當我這樣做時,不會創建一個新張量;而是會返回一個基於底層數據的不同域段(view)的張量。這意味著,如果我編輯該視角下的這些數據,它就會反映在原始的張量中。

在這種情況下,了解如何做到這一點並不算太困難:3 和 4 位於鄰接的內存中,我們只需要記錄一個說明該(邏輯)張量的數據位於頂部以下 2 個位置的偏移量(offset)。(每個張量都記錄一個偏移量,但大多數時候它為零,出現這種情況時我會在我的圖表中省略它。)


演講時的提問:如果我取張量的一個域段,我該如何釋放底層張量的內存?

答案:你必須製作該域段的一個副本,由此斷開其與原始物理內存的連接。你能做的其它事情實際上並不多。另外,如果你很久之前寫過 Java,取一個字元串的子字元串也有類似的問題,因為默認不會製作副本,所以子字元串會保留(可能非常大的字元串)。很顯然,Java 7u6 將其固定了下來。

如果我想取第一列,還會更有意思:

萬字綜述,核心開發者全面解讀PyTorch內部機制

當我們查看物理內存時,可以看到該列的元素不是相鄰的:兩者之間有一個元素的間隙。步幅在這裡就大顯神威了:我們不再將一個元素與下一個元素之間的步幅指定為 1,而是將其設定為 2,即跳兩步。(順便一提,這就是其被稱為「步幅(stride)」的原因:如果我們將索引看作是在布局上行走,步幅就指定了我們每次邁步時向前多少位置。)

步幅表示實際上可以讓你表示所有類型的張量域段;如果你想了解各種不同的可能做法,請參閱 https://ezyang.github.io/stride-visualizer/index.html

我們現在退一步看看,想想我們究竟如何實現這種功能(畢竟這是一個關於內部機制的演講)。如果我們可以得到張量的域段,這就意味著我們必須解耦張量的概念(你所知道且喜愛的面向用戶的概念)以及存儲張量的數據的實際物理數據的概念(稱為「存儲(storage)」):

萬字綜述,核心開發者全面解讀PyTorch內部機制

也許會有多個張量共享同一存儲。存儲會定義張量的 dtype 和物理大小,同時每個張量還會記錄大小、步幅和偏移量,這定義的是物理內存的邏輯解釋。

有一點需要注意:總是會存在一個張量-存儲對,即使並不真正需要存儲的「簡單」情況也是如此(比如,只是用 torch.zeros(2, 2) 劃配一個鄰接張量時)。


順便一提,我們感興趣的不是這種情況,而是有一個分立的存儲概念的情況,只是將一個域段定義為有一個基張量支持的張量。這會更加複雜一些,但也有好處:鄰接張量可以實現遠遠更加直接的表示,而沒有存儲造成的間接麻煩。這樣的變化能讓 PyTorch 的內部表示方式更接近 Numpy。

我們已經介紹了一些張量的數據布局(有人可能會說,如果你正確地理解了數據表示,其它一切都會自然到位)。但還是有必要簡要談談如何實現對張量的操作。在最抽象的層面上,當你調用 torch.mm 時,會發生兩次調度:

萬字綜述,核心開發者全面解讀PyTorch內部機制

第一次調度基於設備類型和張量布局:比如是 CPU 張量還是 CUDA張量,是有步幅的張量還是稀疏的張量。這個調度是動態的:這是一個虛函數(virtual function)調用(這個虛函數調用究竟發生在何處是本演講後半部分的主題)。這裡需要做一次調度應該是合理的:CPU 矩陣乘法的實現非常不同於 CUDA 的實現。這裡是動態調度的原因是這些核(kernel)可能位於不同的庫(比如 libcaffe2.so 或 libcaffe2_gpu.so),這樣你就別無選擇:如果你想進入一個你沒有直接依賴的庫,你必須通過動態調度抵達那裡。

第二次調度是在所涉 dtype 上的調度。這個調度只是一個簡單的 switch 語句,針對的是核選擇支持的任意 dtype。這裡需要調度的原因也很合理:CPU 代碼(或 CUDA 代碼)是基於 float 實現乘法,這不同於用於 int 的代碼。這說明你需要為每種 dtype 都使用不同的核。

如果你想要理解 PyTorch 中運算元的調用方式,這可能就是你頭腦中應有的最重要的知識。後面當我們更深入代碼時還會回到這裡。

萬字綜述,核心開發者全面解讀PyTorch內部機制

因為我們已經談過了張量,所以我還想花點時間談談張量擴展。畢竟,除了密集的 CPU 浮點數張量,還有其它很多類型的張量,比如 XLA 張量、量化張量、MKL-DNN 張量;而對於一個張量庫,還有一件需要思考的事情:如何兼顧這些擴展?

萬字綜述,核心開發者全面解讀PyTorch內部機制

我們當前的用於擴展的模型提供了張量的四個擴展點。首先,有三個獨立地確定張量類型的配套參數:

  • device(設備):描述了實際存儲張量的物理內存,比如在 CPU、英偉達 GPU(cuda)、AMD GPU(hip)或 TPU(xla)上。設備之間各不相同的特性是有各自自己的分配器(allocator),這沒法用於其它設備。
  • layout(布局):描述了對物理內存進行邏輯解讀的方式。最常用的布局是有步幅的張量(strided tensor),但稀疏張量的布局不同,其涉及到一對張量,一個用於索引,一個用於數據;MKL-DNN 張量的布局更加奇特,比如 blocked layout,僅用步幅不能表示它。
  • dtype(數據類型):描述了張量中每個元素實際存儲的數據的類型,比如可以是浮點數、整型數或量化的整型數。

如果你想為 PyTorch 張量添加一種擴展,你應該思考你想要擴展這些參數中的哪幾種。這些參數的笛卡爾積定義了你可以得到的所有可能的張量。現在,並非所有這些組合都有核(誰為 FPGA 上的稀疏量化張量用核?),但原則上這種組合可能有意義,因此我們至少應該支持表達它。

要為張量的功能添加「擴展」,還有最後一種方法,即圍繞能實現的目標類型的 PyTorch 張量編寫一個 wrapper(包裝)類。這可能聽起來理所當然,但有時候人們在只需要製作一個 wrapper 類時卻跑去擴展那三個參數。wrapper 類的一個突出優點是開發結果可以完全不影響原來的類型(out of tree)。

你何時應該編寫張量 wrapper,而不是擴展 PyTorch 本身?關鍵的指標是你是否需要將這個張量傳遞通過 autograd(自動梯度)反向通過過程。舉個例子,這個指標告訴我們稀疏張量應該是一種真正的張量擴展,而不只是一種包含一個索引和值張量的 Python 對象:當在涉及嵌入的網路上執行優化時,我們想要嵌入生成稀疏的梯度。

萬字綜述,核心開發者全面解讀PyTorch內部機制

我們對擴展的理念也會影響張量本身的數據布局。對於我們的張量結構,我們真正想要的一件事物是固定的布局:我們不想要基本操作(這個說法很常見),比如「一個張量的大小是多少?」來請求虛調度。

所以當你查看一個張量的實際布局時(定義為 TensorImpl 結構),會看到所有欄位的一個公共前綴——我們認為所有類似「張量」的東西都會有;還有一些欄位僅真正適用於有步幅的張量,但它們也很重要,所以我們將其保留在主結構中;然後可以在每個張量的基礎上完成有自定義欄位的後綴。比如稀疏張量可將其索引和值存儲在這個後綴中。

自動梯度(autograd)

我已經說明了張量,但如果 PyTorch 僅有這點把戲,這就只不過是 Numpy 的克隆罷了。PyTorch 的顯著特性是其在最初發布時就已提供對張量的自動微分(現在我們還有 TorchScript 等炫酷功能,但那時候就只有這個!)

自動微分是做啥?這是負責運行神經網路的機制:

萬字綜述,核心開發者全面解讀PyTorch內部機制

……以及填充實際計算你的網路的梯度時所缺少的代碼:

萬字綜述,核心開發者全面解讀PyTorch內部機制

花點時間看看這幅圖。其中有很多東西需要解讀,我們來看看:

  • 首先將你的目光投向紅色和藍色的變數。PyTorch 實現了反向模式自動微分,這意味著我們可以「反向」走過前向計算來有效地計算梯度。查看變數名就能看到這一點:在紅色部分的底部,我們計算的是損失(loss);然後在這個程序的藍色部分,我們所做的第一件事是計算 grad_loss。loss 根據 next_h2 計算,這樣我們可以計算出 grad_next_h2。從技術上講,我們加了 grad_ 的變數其實並不是梯度,它們實際上左乘了一個向量的雅可比矩陣,但在 PyTorch 中,我們就稱之為 grad,基本上所有人都知道這是什麼意思。
  • 如果代碼的結構保持一樣,而行為沒有保持一樣:來自前向的每一行都被替換為一個不同的計算,其代表了前向運算的導數。舉個例子,tanh 運算被轉譯成了 tanh_backward 運算(這兩行用圖左邊一條灰線連接)。前向和反向運算的輸入和輸出交換:如果前向運算得到 next_h2,反向運算就以 grad_next_h2 為輸入。

autograd 的意義就在於執行這幅圖所描述的計算,但卻不用真正生成這個源。PyTorch autograd 並不執行源到源的變換(儘管 PyTorch JIT 確實知道如何執行符號微分(symbolic differentiation))。

萬字綜述,核心開發者全面解讀PyTorch內部機制

要做到這一點,我們需要在張量上執行運算時存儲更多元數據。讓我們調整一下我們對張量數據結構的圖:現在不只是一個指向存儲的張量,我們還有一個包裝這個張量的變數,而且也存儲更多信息(AutogradMeta),這是用戶在自己的 PyTorch 腳本中調用 loss.backward() 執行 autograd 時所需的。

這張幻燈片的內容在不久的將來就會過時。Will Feng 在簡單融合了 PyTorch 的前端埠之後,正在推動 C++ 中變數和張量的融合:https://github.com/pytorch/pytorch/issues/13638。

我們也必須更新上面關於調度的圖:

萬字綜述,核心開發者全面解讀PyTorch內部機制

在我們調度到 CPU 或 CUDA 實現之前,還有另一個對變數的調度,其負責打開(unwrap)變數,調用底層實現(綠色),然後再重新將結果包裝進變數並為反向過程記錄必需的 autograd 元數據。

某些實現不會 unwrap;它們只是調用其它變數實現。所以你可能要在變數宇宙中花些時間。但是,一旦你 unwrap 並進入了非變數張量宇宙,你就到達終點了;你再也不用退回變數(除非從你的函數返回)。

在我的紐約聚會演講中,我跳過了以下七頁幻燈片。對它們的文本介紹還要等一段時間。

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

萬字綜述,核心開發者全面解讀PyTorch內部機制

工程開發

說夠了概念,我們來看看代碼。

找到你的路徑

PyTorch 有大量文件夾,在 CONTRIBUTING.md 文檔中有對它們的非常詳細的描述,但實際上你只需知曉 4 個目錄:

萬字綜述,核心開發者全面解讀PyTorch內部機制

  • 首先,torch/ 包含你最熟悉的東西:你導入和使用的實際的 Python 模塊。這些東西是 Python 代碼而且易於操作(只需要進行修改然後查看結果即可)。但是,如果太過深入……
  • torch/csrc/:實現了你可能稱為 PyTorch 前端的 C++ 代碼。用更描述性的術語講,它實現了在 Python 和 C++ 間轉換的綁定代碼(binding code);另外還有一些相當重要的 PyTorch 部分,比如 autograd 引擎和 JIT 編譯器。它也包含 C++ 前端代碼。
  • aten/:這是「A Tensor Library」的縮寫(由 Zachary DeVito 命名),是一個實現張量運算的 C++ 庫。如果你檢查某些核代碼所處的位置,很可能就在 ATen。ATen 本身就分為兩個運算元區域:「原生」運算元(運算元的現代的 C++ 實現)和「傳統」運算元(TH、THC、THNN、THCUNN),這些是遺留的 C 實現。傳統的運算元是其中糟糕的部分;如果可以,請勿在上面耗費太多時間。
  • c10/:這是「Caffe2」和「A"Ten"」的雙關語,包含 PyTorch 的核心抽象,包括張量和存儲數據結構的實際實現。

找代碼需要看很多地方;我們應該簡化目錄結構,就是這樣。如果你想研究運算元,你應該在 aten 上花時間。

我們看看在實踐中是如何分離這些代碼的:

萬字綜述,核心開發者全面解讀PyTorch內部機制

當你調用一個函數時,比如 torch.add,會發生什麼?如果你記得我們的有關調度的討論,你腦中應該已有了這些基礎:

  • 我們必須從 Python 國度轉換到 C++ 國度(Python 參數解析)。
  • 我們處理變數調度(VariableType—Type,順便一提,和編程語言類型並無特別關聯,只是一個用於執行調度的小工具)。
  • 我們處理設備類型/布局調度(Type)。
  • 我們有實際的核,這要麼是一個現代的原生函數,要麼是傳統的 TH 函數。

其中每一步都具體對應於一些代碼。讓我們開路穿過這片叢林。

萬字綜述,核心開發者全面解讀PyTorch內部機制

我們在 C++ 代碼中的起始著陸點是一個 Python 函數的 C 實現,我們已經在 Python 那邊見過它,像是 torch._C.VariableFunctions.add。THPVariable_add 就是這樣一個實現。

對於這些代碼,有一點很重要:這些代碼是自動生成的。如果你在 GitHub 庫中搜索,你沒法找到它們,因為你必須實際 build PyTorch 才能看到它們。另外一點也很重要:你不需要真正深入理解這些代碼是在做什麼,你應該快速瀏覽它,知道它的功能。

我在上面用藍色標註了最重要的部分:你可以看到這裡使用了一個 PythonArgParser 類來從 Python args 和 kwargs 取出 C++ 對象;然後我們調用一個 dispatch_add 函數(紅色內聯);這會釋放全局解釋器鎖,然後調用在 C++ 張量自身上的一個普通的舊方法。在其回來的路上,我們將返回的 Tensor 重新包裝進 PyObject。

(這裡幻燈片中有個錯誤:我應該講解變數調度代碼。我這裡還沒有修復。某些神奇的事發生了,於是……)

萬字綜述,核心開發者全面解讀PyTorch內部機制

當我們在 Tensor 類上調用 add 方法時,還沒有虛調度發生。相反,我有一個內聯方法,其調用了一個內聯方法,其會在「Type」對象上調用一個虛方法。這個方法是真正的虛方法(這就是我說 Type 只是一個讓你實現動態調度的「小工具」的原因)。

在這個特定案例中,這個虛調用會調度到在一個名為 TypeDefault 的類上的 add 的實現。這剛好是因為我們有一個對所有設備類型(CPU 和 CUDA)都一樣的 add 的實現;如果我們剛好有不同的實現,我們可能最終會得到 CPUFloatType::add 這樣的結果。正是這種虛方法的實現能讓我們最終得到實際的核代碼。

也希望這張幻燈片很快過時;Roy Li 正在研究使用另一種機制替代 Type 調度,這能讓我們更好地在移動端上支持 PyTorch。

值得再次強調,一直到我們到達核,所有這些代碼都是自動生成的。

萬字綜述,核心開發者全面解讀PyTorch內部機制

道路蜿蜒曲折,一旦你能基本上把握方向了,我建議你直接跳到核部分。

編寫核(kernel)

PyTorch 為有望編寫核的人提供了大量有用工具。在這一節我們會了解其中一些。但首先,編寫核需要什麼?

萬字綜述,核心開發者全面解讀PyTorch內部機制

我們一般將 PyTorch 中的核看作由以下部分組成:

  • 首先有一些我們要寫的有關核的元數據,這能助力代碼生成並讓你獲取所有與 Python 的捆綁包,同時無需寫任何一行代碼。
  • 一旦你到達了核,你就經過了設備類型/布局調度。你首先需要寫的是錯誤檢查,以確保輸入的張量有正確的維度。(錯誤檢查真正很重要!不要吝惜它!)
  • 接下來,我們一般必須分配我們將要寫入輸出的結果張量。
  • 該到寫核的時候了。現在你應該做第二次 dtype 調度,以跳至其所操作的每個 dtype 特定的核。(你不應該過早做這件事,因為那樣的話你就會毫無用處地複製在任何情況下看起來都一樣的代碼。)
  • 大多數高性能核都需要某種形式的並行化,這樣就能利用多 CPU 系統了。(CUDA 核是「隱式」並行化的,因為它們的編程模型構建於大規模並行化之上。)
  • 最後,你需要讀取數據並執行你想做的計算!

在後面的幻燈片中,我將介紹 PyTorch 中能幫你實現這些步驟的工具。

萬字綜述,核心開發者全面解讀PyTorch內部機制

要充分利用 PyTorch 的代碼生成能力,你需要為你的運算元寫一個模式(schema)。這個模式能提供你的函數的 mypy 風格類型,並控制是否為 Tensor 上的方法或函數生成捆綁包。你還可以告訴模式針對給定的設備-布局組合,應該調用你的運算元的哪種實現。

有關這種格式的更多信息,請參閱:https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/README.md

萬字綜述,核心開發者全面解讀PyTorch內部機制

你可能也需要為你在 derivatives.yaml 中的操作定義一個導數。

萬字綜述,核心開發者全面解讀PyTorch內部機制

錯誤檢查可以在低層 API 完成,也能通過高層 API 實現。低層 API 只是一個宏 TORCH_CHECK,其接收的是一個布爾值,然後還有任意數量的參數構成錯誤字元串(error string)以便得出結論看該布爾值是否為真。

這個宏有個很好的地方:你可以將字元串與非字元串數據混合起來;每一項都使用它們的 operator<< 實現進行格式化,PyTorch 中大多數重要的數據類型都有 operator<< 實現。

高層 API 能讓你免於反覆編寫重複的錯誤消息。其工作方法是;你首先將每個張量包裝為 TensorArg,這包含有關張量來處的信息(比如其參數名稱)。然後它提供了一些預先裝好的用於檢查多種屬性的函數;比如 checkDim() 測試的是張量的維度是否是一個固定數值。如果不是,該函數就基於 TensorArg 元數據提供一個用戶友好的錯誤消息。

萬字綜述,核心開發者全面解讀PyTorch內部機制

在用 PyTorch 寫運算元時,有一點很重要:你往往要註冊三個運算元:abs_out(其操作的是一個預分配的輸出,其實現了 out= keyword 參數)、abs_(其操作的是 inplace)、abs(這只是一個運算元的普通的舊功能版本)。

大部分時間,abs_out 是真正的主力,abs 和 abs_ 只是圍繞 abs_out 的薄弱 wrapper;但有時候也可為每個案例編寫專門的實現。

萬字綜述,核心開發者全面解讀PyTorch內部機制

要執行 dtype 調度,你應該使用 AT_DISPATCH_ALL_TYPES 宏。這會獲取你想要進行調度操作的張量的 dtype,並還會為可從該宏調度的每個 dtype 指定一個 lambda。通常而言,這個 lambda 只是調用一個模板輔助函數。

這個宏不只是「執行調度」,它也會決定你的核將支持的 dtype。這樣,這個宏實際上就有相當多一些版本,這能讓你選取不同的 dtype 子集以生成特定結果。大多數時候,你只需要 AT_DISPATCH_ALL_TYPES,但也要關注你可能需要調度其它更多類型的情況。

萬字綜述,核心開發者全面解讀PyTorch內部機制

在 CPU 上,你通常需要並行化你的代碼。過去,這通常是通過直接在你的代碼中添加 OpenMP pragma 來實現。

萬字綜述,核心開發者全面解讀PyTorch內部機制

某些時候,你必須真正訪問數據。PyTorch 為此提供了相當多一些選擇。

  • 如果你只想獲取某個特定位置的值,你應該使用 TensorAccessor。張量存取器就像是一個張量,但它將張量的維度和 dtype 硬編碼為了模板參數。當你檢索一個存取器時,比如 x.accessor
  • ();,我們會做一次運行時間測試以確保張量確實是這種格式;但那之後,每次存取都不會被檢查。張量存取器能正確地處理步幅,因此你最好使用它們,而不是原始的指針訪問(不幸的是,很多傳統的核是這樣做的)。另外還有 PackedTensorAccessor,這特別適用於通過 CUDA launch 發送存取器,這樣你就能從你的 CUDA 核內部獲取存取器。(一個值得一提的問題:TensorAccessor 默認是 64 位索引,這比 CUDA 中的 32 位索引要慢得多!)
  • 如果你在用很常規的元素存取編寫某種運算元,比如逐點運算,那麼使用遠遠更高級的抽象要好得多,比如 TensorIterator。這個輔助類能為你自動處理廣播和類型提升(type promotion),相當好用。
  • 要在 CPU 上獲得真正的速度,你可能需要使用向量化的 CPU 指令編寫你的核。我們也有用於這方面的輔助函數!Vec256 類表示一種標量向量,並提供了一些能在它們上一次性執行向量化運算的方法。然後 binary_kernel_vec 等輔助函數能讓你輕鬆地運行向量化運算,然後結束那些沒法用普通的舊指令很好地轉換成向量指令的東西。這裡的基礎設施還能在不同指令集下多次編譯你的核,然後在運行時間測試你的 CPU 支持什麼指令,再在這些情況中使用最佳的核。

萬字綜述,核心開發者全面解讀PyTorch內部機制

PyTorch 中大量核都仍然是用傳統的 TH 風格編寫的。(順便一提,TH 代表 TorcH。這是個很好的縮寫詞,但很不幸被污染了;如果你看到名稱中有 TH,可認為它是傳統的。)傳統 TH 風格是什麼意思呢?

  • 它是以 C 風格書寫的,沒有(或很少)使用 C++。
  • 其 refcounted 是人工的(使用了對 THTensor_free 的人工調用以降低你使用張量結束時的 refcounts)。
  • 其位於 generic/ 目錄,這意味著我們實際上要編譯這個文件很多次,但要使用不同的 #define scalar_t

這種代碼相當瘋狂,而且我們討厭回顧它,所以請不要添加它。如果你想寫代碼但對核編寫了解不多,你能做的一件有用的事情:將某些 TH 函數移植到 ATen。

工作流程效率

萬字綜述,核心開發者全面解讀PyTorch內部機制

最後我想談談在 PyTorch 上的工作效率。如果 PyTorch 那龐大的 C++ 代碼庫是阻攔人們為 PyTorch 做貢獻的第一隻攔路虎,那麼你的工作流程的效率就是第二隻。如果你想用 Python 習慣開發 C++,那可能會很艱辛:重新編譯 PyTorch 需要大量時間,你也需要大量時間才能知道你的修改是否有效。

如何高效工作本身可能就值得做一場演講,但這頁幻燈片總結了一些我曾見過某些人抱怨的最常見的反模式:「開發 PyTorch 很困難。」

  • 如果你編輯一個 header,尤其是被許多源文件包含的 header(尤其當被 CUDA 文件包含時),可以預見會有很長的重新 build 時間。盡量只編輯 cpp 文件,編輯 header 要審慎!
  • 我們的 CI 是一種非常好的零設置的測試修改是否有效的方法。但在獲得返回信號之前你可能需要等上一兩個小時。如果你在進行一種將需要大量實驗的改變,那就花點時間設置一個本地開發環境。類似地,如果你在特定的 CI 配置上遇到了困難的 debug 問題,就在本地設置它。你可以將 Docker 鏡像下載到本地並運行:https://github.com/pytorch/ossci-job-dsl
  • 貢獻指南解釋了如何設置 ccache:https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#use-ccache ;強烈建議這個,因為這可以讓你在編輯 header 時幸運地避免大量重新編譯。當我們在不應該重新編譯文件時重新編譯時,這也能幫你覆蓋我們的 build 系統的漏洞。
  • 最後,我們會有大量 C++ 代碼。如果你是在一台有 CPU 和 RAM 的強大伺服器上 build,那麼會有很愉快的體驗。特別要說明,我不建議在筆記本電腦上執行 CUDA build。build CUDA 非常非常慢,而筆記本電腦往往性能不足,不足以快速完成。

參與進來!

萬字綜述,核心開發者全面解讀PyTorch內部機制

這就是我們旋風一般的 PyTorch 內核之旅了!其中省略了很多很多東西;但希望這裡的描述和解釋至少能幫你消化其代碼庫中相當大一部分。

接下來該做什麼?你能做出怎樣的貢獻?我們的問題跟蹤器是個開始的好地方:https://github.com/pytorch/pytorch/issues。

從今年開始,我們一直在分類鑒別問題;標註有「triaged」的問題表示至少有一個 PyTorch 開發者研究過它並對該問題進行了初步評估。你可以使用這些標籤找到我們認為哪些問題是高優先順序的或查看針對特定模塊(如 autograd)的問題,也能找到我們認為是小問題的問題。(警告:我們有時是錯的!)

即使你並不想馬上就開始寫代碼,也仍有很多其它有用的工作值得去做,比如改善文檔(我很喜歡合併文檔 PR,它們都很贊)、幫助我們重現來自其他用戶的 bug 報告以及幫助我們討論問題跟蹤器上的 RFC。沒有我們的開源貢獻者,PyTorch 不會走到今天;我們希望你也能加入我們!

原文地址:http://blog.ezyang.com/2019/05/pytorch-internals/

喜歡這篇文章嗎?立刻分享出去讓更多人知道吧!

本站內容充實豐富,博大精深,小編精選每日熱門資訊,隨時更新,點擊「搶先收到最新資訊」瀏覽吧!


請您繼續閱讀更多來自 機器之心 的精彩文章:

如何引誘分類器犯錯?南大周志華等提出用自編碼器生成惡意訓練數據
推理速度提升5.1倍,參數減少88%:谷歌提出新型卷積網路EfficientNet

TAG:機器之心 |