pytorch內部機制解讀

斯坦福大學博士生與 Facebook 人工智能研究所研究工程師 Edward Z. Yang 是 PyTorch 開源項目的核心開發者之一。本文是對Mr. Yang寫的關於pytorch內部機制的博文的總結,該文編寫時pytorch版本估計爲0.3,此時pytorch的tensor和Viarable還沒完全合併。

原文地址:http://blog.ezyang.com/2019/05/pytorch-internals/
中文翻譯來源:https://www.xumingchao.com/pytorch%E5%86%85%E9%83%A8%E6%9C%BA%E5%88%B6%E8%A7%A3%E8%AF%BB-%E7%BB%BC%E8%BF%B0/
在這裏插入圖片描述
本演講的目的是爲你提供一份導航圖:爲你講解一個「支持自動微分的Tensor庫」的基本概念結構,併爲你提供一些能幫你在代碼庫中尋路的工具和技巧。我預設你之前已經寫過一些 PyTorch,但卻可能還沒有深入理解機器學習軟件庫的編寫方式。
在這裏插入圖片描述
本篇文章分爲兩個部分介紹pytorch的內部機制:第一部分是全面介紹張量庫的各種概念,詳細討論這種數據類型究竟能提供什麼,這能讓我們更好地理解其內部真正的實現方式。

我們也會談到「擴展點(extension points)」的三個概念、佈局(layout)、設備(device)和數據類型(dtype),這能引導我們思考Tensor類的擴展的方式。同時也會對自動梯度(autograd)進行講解。
第二部分會闡述真正用 PyTorch 寫代碼時所涉及的基本細節。我會告訴你如何在 autograd 代碼中披荊斬棘、什麼代碼是真正重要的以及怎樣造福他人,我還會介紹 PyTorch 爲你寫核(kernel)所提供的所有炫酷工具。

概念

張量 Tensor

Tensor是 PyTorch 中的核心數據結構。對於Tensor直觀上所表示的東西,你可能已有很好的理解:Tensor是一種包含某種標量類型(比如浮點數和整型數等)的 n 維數據結構。我們可以將Tensor看作是由一些數據構成的,還有一些元數據描述了Tensor的大小、所包含的元素的類型(dtype)、Tensor所在的設備(CPU 內存?CUDA 內存?)
在這裏插入圖片描述
另外還有一個你可能沒那麼熟悉的元數據:步幅(stride)。stride 實際上是 PyTorch 最別緻的特徵之一,所以值得稍微多討論它一些。
在這裏插入圖片描述
Tensor一個數學概念。但要在我們的計算機中表示它,我們必須爲它們定義某種物理表示方法。最常用的表示方法是在內存中相鄰地放置Tensor的每個元素(這也是術語「contiguous(鄰接)」的來源),即將每一行寫出到內存,如上所示。在上面的案例中,我已經指定該Tensor包含 32 位的整型數,這樣你可以看到每一個整型數都位於一個物理地址中,每個地址與相鄰地址相距 4 字節。爲了記住Tensor的實際維度,我們必須將規模大小記爲額外的元數據。

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

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

stride是我們爲 PyTorch 用戶講解方法的基本基礎。舉個例子,假設我想取出一個表示以上Tensor的第二行的Tensor:
在這裏插入圖片描述
在這種情況下,瞭解如何做到這一點並不算太困難:3 和 4 位於鄰接的內存中,我們只需要記錄一個說明該(邏輯)Tensor的數據位於頂部以下 2 個位置的偏移量(offset)。(每個Tensor都記錄一個偏移量,但大多數時候它爲零,出現這種情況時我會在我的圖表中省略它。)

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

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

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

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

我們現在退一步看看,想想我們究竟如何實現這種功能(畢竟這是一個關於內部機制的演講)。如果我們可以得到Tensor的域段,這就意味着我們必須解耦Tensor的概念(你所知道且喜愛的面向用戶的概念)以及存儲Tensor的數據的實際物理數據的概念(稱爲「存儲(storage)」):
在這裏插入圖片描述
也許會有多個Tensor共享同一存儲。存儲會定義Tensor的 dtype 和物理大小,同時每個Tensor還會記錄大小、stride和偏移量,這定義的是物理內存的邏輯解釋。

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

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

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

在這裏插入圖片描述
第一次調度基於設備類型和Tensor佈局:比如是 CPU Tensor還是 CUDATensor,是有stride的Tensor還是稀疏的Tensor。這個調度是動態的:這是一個虛函數(virtual function)調用(這個虛函數調用究竟發生在何處是本演講後半部分的主題)。
這裏需要做一次調度應該是合理的:CPU 矩陣乘法的實現非常不同於 CUDA 的實現。這裏是動態調度的原因是這些核(kernel)可能位於不同的庫(比如 libcaffe2.so 或 libcaffe2_gpu.so),這樣你就別無選擇:如果你想進入一個你沒有直接依賴的庫,你必須通過動態調度抵達那裏。
第二次調度是在所涉 dtype 上的調度。這個調度只是一個簡單的 switch 語句,針對的是核選擇支持的任意 dtype。這裏需要調度的原因也很合理:CPU 代碼(或 CUDA 代碼)是基於 float 實現乘法,這不同於用於 int 的代碼。這說明你需要爲每種 dtype 都使用不同的核。

如果你想要理解 PyTorch 中算子的調用方式,這可能就是你頭腦中應有的最重要的知識。後面當我們更深入代碼時還會回到這裏。
在這裏插入圖片描述
因爲我們已經談過了Tensor,所以我還想花點時間談談Tensor擴展。畢竟,除了密集的 CPU 浮點數Tensor,還有其它很多類型的Tensor,比如 XLA Tensor、量化Tensor、MKL-DNN Tensor;而對於一個Tensor庫,還有一件需要思考的事情:如何兼顧這些擴展?
在這裏插入圖片描述
我們當前的用於擴展的模型提供了Tensor的四個擴展點。首先,有三個獨立地確定Tensor類型的配套參數:

  • device(設備):描述了實際存儲Tensor的物理內存,比如在 CPU、英偉達 GPU(cuda)、AMD GPU(hip)或 TPU(xla)上。設備之間各不相同的特性是有各自自己的分配器(allocator),這沒法用於其它設備。

  • layout(佈局):描述了對物理內存進行邏輯解讀的方式。最常用的佈局是有stride的Tensor(strided tensor),但稀疏Tensor的佈局不同,其涉及到一對Tensor,一個用於索引,一個用於數據;MKL-DNN Tensor的佈局更加奇特,比如 blocked layout,僅用stride不能表示它。

  • dtype(數據類型):描述了Tensor中每個元素實際存儲的數據的類型,比如可以是浮點數、整型數或量化的整型數。

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

要爲Tensor的功能添加「擴展」,還有最後一種方法,即圍繞能實現的目標類型的 PyTorch Tensor編寫一個 wrapper(包裝)類。這可能聽起來理所當然,但有時候人們在只需要製作一個 wrapper 類時卻跑去擴展那三個參數。wrapper 類的一個突出優點是開發結果可以完全不影響原來的類型(out of tree)。
你何時應該編寫Tensor wrapper,而不是擴展 PyTorch 本身?關鍵的指標是你是否需要將這個Tensor傳遞通過 autograd(自動梯度)反向通過過程。舉個例子,這個指標告訴我們稀疏Tensor應該是一種真正的Tensor擴展,而不只是一種包含一個索引和值Tensor的 Python 對象:當在涉及嵌入的網絡上執行優化時,我們想要嵌入生成稀疏的梯度。
在這裏插入圖片描述
我們對擴展的理念也會影響Tensor本身的數據佈局。對於我們的Tensor結構,我們真正想要的一件事物是固定的佈局:我們不想要基本操作(這個說法很常見),比如「一個Tensor的大小是多少?」來請求虛調度。

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

自動梯度(autograd)

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

自動微分是做啥?這是負責運行神經網絡的機制:
在這裏插入圖片描述
……以及填充實際計算你的網絡的梯度時所缺少的代碼:
在這裏插入圖片描述
花點時間看看這幅圖。其中有很多東西需要解讀,我們來看看:

  • 首先將你的目光投向紅色和藍色的變量。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))。
在這裏插入圖片描述
要做到這一點,我們需要在Tensor上執行運算時存儲更多元數據。讓我們調整一下我們對Tensor數據結構的圖:現在不只是一個指向存儲的Tensor,我們還有一個包裝這個Tensor的變量,而且也存儲更多信息(AutogradMeta),這是用戶在自己的 PyTorch 腳本中調用 loss.backward() 執行 autograd 時所需的。

這張幻燈片的內容在不久的將來就會過時。Will Feng 在簡單融合了 PyTorch 的前端端口之後,正在推動 C++ 中變量和Tensor的融合:https://github.com/pytorch/pytorch/issues/13638。
我們也必須更新上面關於調度的圖:
在這裏插入圖片描述
在我們調度到 CPU 或 CUDA 實現之前,還有另一個對變量的調度,其負責打開(unwrap)變量,調用底層實現(綠色),然後再重新將結果包裝進變量併爲反向過程記錄必需的 autograd 元數據。

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

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

工程開發

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

找到你的路徑

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

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

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

我們看看在實踐中是如何分離這些代碼的:
在這裏插入圖片描述
當你調用一個函數時,比如 torch.add,會發生什麼?如果你記得我們的有關調度的討論,你腦中應該已有了這些基礎:

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

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

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

對於這些代碼,有一點很重要:這些代碼是自動生成的。如果你在 GitHub 庫中搜索,你沒法找到它們,因爲你必須實際 build PyTorch 才能看到它們。另外一點也很重要:你不需要真正深入理解這些代碼是在做什麼,你應該快速瀏覽它,知道它的功能。
我在上面用藍色標註了最重要的部分:你可以看到這裏使用了一個 PythonArgParser 類來從 Python args 和 kwargs 取出 C++ 對象;然後我們調用一個 dispatch_add 函數(紅色內聯);這會釋放全局解釋器鎖,然後調用在 C++ Tensor自身上的一個普通的舊方法。在其回來的路上,我們將返回的 Tensor 重新包裝進 PyObject。

(這裏幻燈片中有個錯誤:我應該講解變量調度代碼。我這裏還沒有修復。某些神奇的事發生了,於是……)
在這裏插入圖片描述
當我們在 Tensor 類上調用 add 方法時,還沒有虛調度發生。相反,我有一個內聯方法,其調用了一個內聯方法,其會在「Type」對象上調用一個虛方法。這個方法是真正的虛方法(這就是我說 Type 只是一個讓你實現動態調度的「小工具」的原因)。

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

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

值得再次強調,一直到我們到達核,所有這些代碼都是自動生成的。
在這裏插入圖片描述
道路蜿蜒曲折,一旦你能基本上把握方向了,我建議你直接跳到核部分。

編寫核(kernel)

PyTorch 爲有望編寫核的人提供了大量有用工具。在這一節我們會了解其中一些。但首先,編寫核需要什麼?
在這裏插入圖片描述
我們一般將 PyTorch 中的核看作由以下部分組成:

  • 首先有一些我們要寫的有關核的元數據,這能助力代碼生成並讓你獲取所有與 Python 的捆綁包,同時無需寫任何一行代碼。
  • 一旦你到達了核,你就經過了設備類型/佈局調度。你首先需要寫的是錯誤檢查,以確保輸入的Tensor有正確的維度。(錯誤檢查真正很重要!不要吝惜它!)
  • 接下來,我們一般必須分配我們將要寫入輸出的結果Tensor。
  • 該到寫核的時候了。現在你應該做第二次 dtype 調度,以跳至其所操作的每個 dtype 特定的核。(你不應該過早做這件事,因爲那樣的話你就會毫無用處地複製在任何情況下看起來都一樣的代碼。)
  • 大多數高性能核都需要某種形式的並行化,這樣就能利用多 CPU 系統了。(CUDA 核是「隱式」並行化的,因爲它們的編程模型構建於大規模並行化之上。)
  • 最後,你需要讀取數據並執行你想做的計算!
  • 在後面的幻燈片中,我將介紹 PyTorch 中能幫你實現這些步驟的工具。
    在這裏插入圖片描述
    要充分利用 PyTorch 的代碼生成能力,你需要爲你的算子寫一個模式(schema)。這個模式能提供你的函數的 mypy 風格類型,並控制是否爲 Tensor 上的方法或函數生成捆綁包。你還可以告訴模式針對給定的設備-佈局組合,應該調用你的算子的哪種實現。

有關這種格式的更多信息,請參閱:

https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/README.md
在這裏插入圖片描述
你可能也需要爲你在 derivatives.yaml 中的操作定義一個導數。
在這裏插入圖片描述
錯誤檢查可以在低層 API 完成,也能通過高層 API 實現。低層 API 只是一個宏 TORCH_CHECK,其接收的是一個布爾值,然後還有任意數量的參數構成錯誤字符串(error string)以便得出結論看該布爾值是否爲真。

這個宏有個很好的地方:你可以將字符串與非字符串數據混合起來;每一項都使用它們的 operator<< 實現進行格式化,PyTorch 中大多數重要的數據類型都有 operator<< 實現。
高層 API 能讓你免於反覆編寫重複的錯誤消息。其工作方法是;你首先將每個Tensor包裝爲 TensorArg,這包含有關Tensor來處的信息(比如其參數名稱)。然後它提供了一些預先裝好的用於檢查多種屬性的函數;比如 checkDim() 測試的是Tensor的維度是否是一個固定數值。如果不是,該函數就基於 TensorArg 元數據提供一個用戶友好的錯誤消息。
在這裏插入圖片描述
在用 PyTorch 寫算子時,有一點很重要:你往往要註冊三個算子:abs_out(其操作的是一個預分配的輸出,其實現了 out= keyword 參數)、abs_(其操作的是 inplace)、abs(這只是一個算子的普通的舊功能版本)。

大部分時間,abs_out 是真正的主力,abs 和 abs_ 只是圍繞 abs_out 的薄弱 wrapper;但有時候也可爲每個案例編寫專門的實現。
在這裏插入圖片描述
要執行 dtype 調度,你應該使用 AT_DISPATCH_ALL_TYPES 宏。這會獲取你想要進行調度操作的Tensor的 dtype,並還會爲可從該宏調度的每個 dtype 指定一個 lambda。通常而言,這個 lambda 只是調用一個模板輔助函數。

這個宏不只是「執行調度」,它也會決定你的核將支持的 dtype。這樣,這個宏實際上就有相當多一些版本,這能讓你選取不同的 dtype 子集以生成特定結果。大多數時候,你只需要 AT_DISPATCH_ALL_TYPES,但也要關注你可能需要調度其它更多類型的情況。
在這裏插入圖片描述
在 CPU 上,你通常需要並行化你的代碼。過去,這通常是通過直接在你的代碼中添加 OpenMP pragma 來實現。
在這裏插入圖片描述
某些時候,你必須真正訪問數據。PyTorch 爲此提供了相當多一些選擇。

  • 如果你只想獲取某個特定位置的值,你應該使用 TensorAccessor。Tensor存取器就像是一個Tensor,但它將Tensor的維度和 dtype 硬編碼爲了模板參數。當你檢索一個存取器時,比如 x.accessor();,我們會做一次運行時間測試以確保Tensor確實是這種格式;但那之後,每次存取都不會被檢查。Tensor存取器能正確地處理stride,因此你最好使用它們,而不是原始的指針訪問(不幸的是,很多傳統的核是這樣做的)。另外還有 PackedTensorAccessor,這特別適用於通過 CUDA launch 發送存取器,這樣你就能從你的 CUDA 核內部獲取存取器。(一個值得一提的問題:TensorAccessor 默認是 64 位索引,這比 CUDA 中的 32 位索引要慢得多!)

  • 如果你在用很常規的元素存取編寫某種算子,比如逐點運算,那麼使用遠遠更高級的抽象要好得多,比如 TensorIterator。這個輔助類能爲你自動處理廣播和類型提升(type promotion),相當好用。

  • 要在 CPU 上獲得真正的速度,你可能需要使用向量化的 CPU 指令編寫你的核。我們也有用於這方面的輔助函數!Vec256 類表示一種標量向量,並提供了一些能在它們上一次性執行向量化運算的方法。然後 binary_kernel_vec 等輔助函數能讓你輕鬆地運行向量化運算,然後結束那些沒法用普通的舊指令很好地轉換成向量指令的東西。這裏的基礎設施還能在不同指令集下多次編譯你的核,然後在運行時間測試你的 CPU 支持什麼指令,再在這些情況中使用最佳的核。
    -在這裏插入圖片描述
    PyTorch 中大量核都仍然是用傳統的 TH 風格編寫的。(順便一提,TH 代表 TorcH。這是個很好的縮寫詞,但很不幸被污染了;如果你看到名稱中有 TH,可認爲它是傳統的。)傳統 TH 風格是什麼意思呢?

  • 它是以 C 風格書寫的,沒有(或很少)使用 C++。

  • 其 refcounted 是人工的(使用了對 THTensor_free 的人工調用以降低你使用Tensor結束時的 refcounts)。

  • 其位於 generic/ 目錄,這意味着我們實際上要編譯這個文件很多次,但要使用不同的 #define scalar_t
    這種代碼相當瘋狂,而且我們討厭回顧它,所以請不要添加它。如果你想寫代碼但對覈編寫了解不多,你能做的一件有用的事情:將某些 TH 函數移植到 ATen。

工作流程效率

在這裏插入圖片描述
最後我想談談在 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 內核之旅了!其中省略了很多很多東西;但希望這裏的描述和解釋至少能幫你消化其代碼庫中相當大一部分。

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

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

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

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章