斯坦福大學博士生與 Facebook 人工智能研究所研究工程師 Edward Z. Yang 是 PyTorch 開源項目的核心開發(fā)者之一。他在 5 月 14 日的 PyTorch 紐約聚會上做了一個有關 PyTorch 內(nèi)部機制的演講鸭津，本文是該演講的長文章版本罢浇。
選自ezyang博客缭召，作者：Edward Z. Yang奈籽，機器之心編譯市咽，參與：panda。
大家好烁焙！今天我想談談 PyTorch 的內(nèi)部機制航邢。
這份演講是為用過 PyTorch并且有心為 PyTorch 做貢獻但卻被 PyTorch 那龐大的 C++ 代碼庫勸退的人提供的。沒必要說謊：PyTorch 代碼庫有時候確實讓人難以招架骄蝇。
本演講的目的是為你提供一份導航圖：為你講解一個「支持自動微分的張量庫」的基本概念結構膳殷，并為你提供一些能幫你在代碼庫中尋路的工具和技巧。我預設你之前已經(jīng)寫過一些 PyTorch乞榨，但卻可能還沒有深入理解機器學習軟件庫的編寫方式秽之。
本演講分為兩部分：在第一部分中，我首先會全面介紹張量庫的各種概念吃既。我首先會談談你們知道且喜愛的張量數(shù)據(jù)類型，并詳細討論這種數(shù)據(jù)類型究竟能提供什么跨细，這能讓我們更好地理解其內(nèi)部真正的實現(xiàn)方式鹦倚。
如果你是一位 PyTorch 高級用戶，你可能已經(jīng)熟悉其中大部分材料了冀惭。我們也會談到「擴展點（extension points）」的三個概念震叙、布局（layout）、設備（device）和數(shù)據(jù)類型（dtype）散休，這能引導我們思考張量類的擴展的方式媒楼。在 PyTorch 紐約聚會的現(xiàn)場演講中，我略過了有關自動梯度（autograd）的幻燈片戚丸，但我在這里會進行一些講解划址。
第二部分會闡述真正用 PyTorch 寫代碼時所涉及的基本細節(jié)。我會告訴你如何在 autograd 代碼中披荊斬棘限府、什么代碼是真正重要的以及怎樣造福他人夺颤，我還會介紹 PyTorch 為你寫核（kernel）所提供的所有炫酷工具。
概念
張量
張量是 PyTorch 中的核心數(shù)據(jù)結構胁勺。對于張量直觀上所表示的東西世澜，你可能已有很好的理解：張量是一種包含某種標量類型（比如浮點數(shù)和整型數(shù)等）的 n 維數(shù)據(jù)結構。我們可以將張量看作是由一些數(shù)據(jù)構成的署穗，還有一些元數(shù)據(jù)描述了張量的大小寥裂、所包含的元素的類型（dtype）嵌洼、張量所在的設備（CPU 內(nèi)存？CUDA 內(nèi)存封恰？）
另外還有一個你可能沒那么熟悉的元數(shù)據(jù)：步幅（stride）麻养。stride 實際上是 PyTorch 最別致的特征之一，所以值得稍微多討論它一些俭驮。
張量一個數(shù)學概念回溺。但要在我們的計算機中表示它，我們必須為它們定義某種物理表示方法混萝。最常用的表示方法是在內(nèi)存中相鄰地放置張量的每個元素（這也是術語「contiguous（鄰接）」的來源）遗遵，即將每一行寫出到內(nèi)存，如上所示逸嘀。在上面的案例中车要，我已經(jīng)指定該張量包含 32 位的整型數(shù)，這樣你可以看到每一個整型數(shù)都位于一個物理地址中崭倘，每個地址與相鄰地址相距 4 字節(jié)翼岁。為了記住張量的實際維度，我們必須將規(guī)模大小記為額外的元數(shù)據(jù)司光。
所以這幅圖與步幅有什么關系琅坡？
假設我想要讀取我的邏輯表示中位置張量 [0,1] 的元素。我該如何將這個邏輯位置轉譯為物理內(nèi)存中的位置残家？步幅能讓我們做到這一點：要找到一個張量中任意元素的位置榆俺，我將每個索引與該維度下各自的步幅相乘，然后將它們?nèi)考拥揭黄鹞牖础Ｔ谏蠄D中茴晋，我用藍色表示第一個維度，用紅色表示第二個維度回窘，以便你了解該步幅計算中的索引和步幅诺擅。進行這個求和后，我得到了 2（零索引的）啡直；實際上烁涌，數(shù)字 3 正是位于這個鄰接數(shù)組的起點以下 2 個位置。
（后面我還會談到 TensorAccessor付枫，這是一個處理索引計算的便利類（convenience class）烹玉。當你使用 TensorAccessor 時，不會再操作原始指針阐滩，這些計算過程已經(jīng)為你隱藏了起來二打。）
步幅是我們?yōu)?PyTorch 用戶講解方法的基本基礎。舉個例子掂榔，假設我想取出一個表示以上張量的第二行的張量：
使用高級的索引支持继效，我只需寫出張量 [1, :] 就能得到這一行症杏。重要的是：當我這樣做時，不會創(chuàng)建一個新張量瑞信；而是會返回一個基于底層數(shù)據(jù)的不同域段（view）的張量厉颤。這意味著，如果我編輯該視角下的這些數(shù)據(jù)凡简，它就會反映在原始的張量中逼友。
在這種情況下，了解如何做到這一點并不算太困難：3 和 4 位于鄰接的內(nèi)存中秤涩，我們只需要記錄一個說明該（邏輯）張量的數(shù)據(jù)位于頂部以下 2 個位置的偏移量（offset）帜乞。（每個張量都記錄一個偏移量，但大多數(shù)時候它為零筐眷，出現(xiàn)這種情況時我會在我的圖表中省略它黎烈。）
演講時的提問：如果我取張量的一個域段，我該如何釋放底層張量的內(nèi)存匀谣？
答案：你必須制作該域段的一個副本照棋，由此斷開其與原始物理內(nèi)存的連接。你能做的其它事情實際上并不多武翎。另外烈炭，如果你很久之前寫過 Java，取一個字符串的子字符串也有類似的問題宝恶，因為默認不會制作副本梳庆，所以子字符串會保留（可能非常大的字符串）。很顯然卑惜，Java 7u6 將其固定了下來。
如果我想取第一列驻售，還會更有意思：
當我們查看物理內(nèi)存時露久，可以看到該列的元素不是相鄰的：兩者之間有一個元素的間隙。步幅在這里就大顯神威了：我們不再將一個元素與下一個元素之間的步幅指定為 1欺栗，而是將其設定為 2毫痕，即跳兩步。（順便一提迟几，這就是其被稱為「步幅（stride）」的原因：如果我們將索引看作是在布局上行走消请，步幅就指定了我們每次邁步時向前多少位置。）
步幅表示實際上可以讓你表示所有類型的張量域段类腮；如果你想了解各種不同的可能做法臊泰，請參閱 https://ezyang.github.io/stride-visualizer/index.html
我們現(xiàn)在退一步看看，想想我們究竟如何實現(xiàn)這種功能（畢竟這是一個關于內(nèi)部機制的演講）蚜枢。如果我們可以得到張量的域段缸逃，這就意味著我們必須解耦張量的概念（你所知道且喜愛的面向用戶的概念）以及存儲張量的數(shù)據(jù)的實際物理數(shù)據(jù)的概念（稱為「存儲（storage）」）：
也許會有多個張量共享同一存儲针饥。存儲會定義張量的 dtype 和物理大小，同時每個張量還會記錄大小需频、步幅和偏移量丁眼，這定義的是物理內(nèi)存的邏輯解釋。
有一點需要注意：總是會存在一個張量-存儲對昭殉，即使并不真正需要存儲的「簡單」情況也是如此（比如苞七，只是用 torch.zeros(2, 2) 劃配一個鄰接張量時）。
順便一提挪丢，我們感興趣的不是這種情況蹂风，而是有一個分立的存儲概念的情況，只是將一個域段定義為有一個基張量支持的張量吃靠。這會更加復雜一些硫眨，但也有好處：鄰接張量可以實現(xiàn)遠遠更加直接的表示，而沒有存儲造成的間接麻煩巢块。這樣的變化能讓 PyTorch 的內(nèi)部表示方式更接近 Numpy礁阁。
我們已經(jīng)介紹了一些張量的數(shù)據(jù)布局（有人可能會說，如果你正確地理解了數(shù)據(jù)表示族奢，其它一切都會自然到位）姥闭。但還是有必要簡要談談如何實現(xiàn)對張量的操作。在最抽象的層面上越走，當你調(diào)用 torch.mm 時棚品，會發(fā)生兩次調(diào)度：
第一次調(diào)度基于設備類型和張量布局：比如是 CPU 張量還是 CUDA張量，是有步幅的張量還是稀疏的張量廊敌。這個調(diào)度是動態(tài)的：這是一個虛函數(shù)（virtual function）調(diào)用（這個虛函數(shù)調(diào)用究竟發(fā)生在何處是本演講后半部分的主題）铜跑。
這里需要做一次調(diào)度應該是合理的：CPU 矩陣乘法的實現(xiàn)非常不同于 CUDA 的實現(xiàn)。這里是動態(tài)調(diào)度的原因是這些核（kernel）可能位于不同的庫（比如 libcaffe2.so 或 libcaffe2_gpu.so）骡澈，這樣你就別無選擇：如果你想進入一個你沒有直接依賴的庫锅纺，你必須通過動態(tài)調(diào)度抵達那里。
第二次調(diào)度是在所涉 dtype 上的調(diào)度肋殴。這個調(diào)度只是一個簡單的 switch 語句囤锉，針對的是核選擇支持的任意 dtype。這里需要調(diào)度的原因也很合理：CPU 代碼（或 CUDA 代碼）是基于 float 實現(xiàn)乘法护锤，這不同于用于 int 的代碼官地。這說明你需要為每種 dtype 都使用不同的核。
如果你想要理解 PyTorch 中算子的調(diào)用方式烙懦，這可能就是你頭腦中應有的最重要的知識驱入。后面當我們更深入代碼時還會回到這里。
因為我們已經(jīng)談過了張量，所以我還想花點時間談談張量擴展沧侥。畢竟可霎，除了密集的 CPU 浮點數(shù)張量，還有其它很多類型的張量宴杀，比如 XLA 張量癣朗、量化張量、MKL-DNN 張量旺罢；而對于一個張量庫旷余，還有一件需要思考的事情：如何兼顧這些擴展？
我們當前的用于擴展的模型提供了張量的四個擴展點扁达。首先正卧，有三個獨立地確定張量類型的配套參數(shù)：
- device（設備）：描述了實際存儲張量的物理內(nèi)存，比如在 CPU跪解、英偉達 GPU（cuda）炉旷、AMD GPU（hip）或 TPU（xla）上。設備之間各不相同的特性是有各自自己的分配器（allocator）叉讥，這沒法用于其它設備窘行。
- layout（布局）：描述了對物理內(nèi)存進行邏輯解讀的方式。最常用的布局是有步幅的張量（strided tensor）图仓，但稀疏張量的布局不同罐盔，其涉及到一對張量，一個用于索引救崔，一個用于數(shù)據(jù)惶看；MKL-DNN 張量的布局更加奇特，比如 blocked layout六孵，僅用步幅不能表示它纬黎。
- dtype（數(shù)據(jù)類型）：描述了張量中每個元素實際存儲的數(shù)據(jù)的類型，比如可以是浮點數(shù)劫窒、整型數(shù)或量化的整型數(shù)莹桅。
如果你想為 PyTorch 張量添加一種擴展，你應該思考你想要擴展這些參數(shù)中的哪幾種烛亦。這些參數(shù)的笛卡爾積定義了你可以得到的所有可能的張量。現(xiàn)在懂拾，并非所有這些組合都有核（誰為 FPGA 上的稀疏量化張量用核?）煤禽，但原則上這種組合可能有意義，因此我們至少應該支持表達它岖赋。
要為張量的功能添加「擴展」檬果，還有最后一種方法，即圍繞能實現(xiàn)的目標類型的 PyTorch 張量編寫一個 wrapper（包裝）類。這可能聽起來理所當然选脊，但有時候人們在只需要制作一個 wrapper 類時卻跑去擴展那三個參數(shù)杭抠。wrapper 類的一個突出優(yōu)點是開發(fā)結果可以完全不影響原來的類型（out of tree）。
你何時應該編寫張量 wrapper恳啥，而不是擴展 PyTorch 本身偏灿？關鍵的指標是你是否需要將這個張量傳遞通過 autograd（自動梯度）反向通過過程。舉個例子钝的，這個指標告訴我們稀疏張量應該是一種真正的張量擴展翁垂，而不只是一種包含一個索引和值張量的 Python 對象：當在涉及嵌入的網(wǎng)絡上執(zhí)行優(yōu)化時，我們想要嵌入生成稀疏的梯度硝桩。
我們對擴展的理念也會影響張量本身的數(shù)據(jù)布局沿猜。對于我們的張量結構，我們真正想要的一件事物是固定的布局：我們不想要基本操作（這個說法很常見）碗脊，比如「一個張量的大小是多少啼肩？」來請求虛調(diào)度。
所以當你查看一個張量的實際布局時（定義為 TensorImpl 結構）衙伶，會看到所有字段的一個公共前綴——我們認為所有類似「張量」的東西都會有祈坠；還有一些字段僅真正適用于有步幅的張量，但它們也很重要痕支，所以我們將其保留在主結構中颁虐；然后可以在每個張量的基礎上完成有自定義字段的后綴。比如稀疏張量可將其索引和值存儲在這個后綴中卧须。
自動梯度（autograd）
我已經(jīng)說明了張量腊尚，但如果 PyTorch 僅有這點把戲甩卓，這就只不過是 Numpy 的克隆罷了。PyTorch 的顯著特性是其在最初發(fā)布時就已提供對張量的自動微分（現(xiàn)在我們還有 TorchScript 等炫酷功能，但那時候就只有這個Ｇ獭）
自動微分是做啥？這是負責運行神經(jīng)網(wǎng)絡的機制：
……以及填充實際計算你的網(wǎng)絡的梯度時所缺少的代碼：
花點時間看看這幅圖雨涛。其中有很多東西需要解讀碧浊，我們來看看：
- 首先將你的目光投向紅色和藍色的變量。PyTorch 實現(xiàn)了反向模式自動微分隘击，這意味著我們可以「反向」走過前向計算來有效地計算梯度侍芝。查看變量名就能看到這一點：在紅色部分的底部，我們計算的是損失（loss）埋同；然后在這個程序的藍色部分州叠，我們所做的第一件事是計算 grad_loss。loss 根據(jù) next_h2 計算凶赁，這樣我們可以計算出 grad_next_h2咧栗。從技術上講逆甜，我們加了 grad_ 的變量其實并不是梯度，它們實際上左乘了一個向量的雅可比矩陣致板，但在 PyTorch 中交煞，我們就稱之為 grad，基本上所有人都知道這是什么意思斟或。
- 如果代碼的結構保持一樣素征，而行為沒有保持一樣：來自前向的每一行都被替換為一個不同的計算，其代表了前向運算的導數(shù)缕粹。舉個例子稚茅，tanh 運算被轉譯成了 tanh_backward 運算（這兩行用圖左邊一條灰線連接）。前向和反向運算的輸入和輸出交換：如果前向運算得到 next_h2平斩，反向運算就以 grad_next_h2 為輸入亚享。
autograd 的意義就在于執(zhí)行這幅圖所描述的計算，但卻不用真正生成這個源绘面。PyTorch autograd 并不執(zhí)行源到源的變換（盡管 PyTorch JIT 確實知道如何執(zhí)行符號微分（symbolic differentiation））欺税。
要做到這一點，我們需要在張量上執(zhí)行運算時存儲更多元數(shù)據(jù)揭璃。讓我們調(diào)整一下我們對張量數(shù)據(jù)結構的圖：現(xiàn)在不只是一個指向存儲的張量晚凿，我們還有一個包裝這個張量的變量，而且也存儲更多信息（AutogradMeta）瘦馍，這是用戶在自己的 PyTorch 腳本中調(diào)用 loss.backward() 執(zhí)行 autograd 時所需的歼秽。
這張幻燈片的內(nèi)容在不久的將來就會過時。Will Feng 在簡單融合了 PyTorch 的前端端口之后情组，正在推動 C++ 中變量和張量的融合：https://github.com/pytorch/pytorch/issues/13638燥筷。
我們也必須更新上面關于調(diào)度的圖：
在我們調(diào)度到 CPU 或 CUDA 實現(xiàn)之前，還有另一個對變量的調(diào)度院崇，其負責打開（unwrap）變量肆氓，調(diào)用底層實現(xiàn)（綠色），然后再重新將結果包裝進變量并為反向過程記錄必需的 autograd 元數(shù)據(jù)底瓣。
某些實現(xiàn)不會 unwrap谢揪；它們只是調(diào)用其它變量實現(xiàn)。所以你可能要在變量宇宙中花些時間捐凭。但是拨扶，一旦你 unwrap 并進入了非變量張量宇宙，你就到達終點了茁肠；你再也不用退回變量（除非從你的函數(shù)返回）屈雄。
在我的紐約聚會演講中，我跳過了以下七頁幻燈片官套。對它們的文本介紹還要等一段時間。
工程開發(fā)
說夠了概念，我們來看看代碼奶赔。
找到你的路徑
PyTorch 有大量文件夾惋嚎，在 CONTRIBUTING.md 文檔中有對它們的非常詳細的描述，但實際上你只需知曉 4 個目錄：
- 首先站刑，torch/ 包含你最熟悉的東西：你導入和使用的實際的 Python 模塊另伍。這些東西是 Python 代碼而且易于操作（只需要進行修改然后查看結果即可）。但是绞旅，如果太過深入……
- torch/csrc/：實現(xiàn)了你可能稱為 PyTorch 前端的 C++ 代碼摆尝。用更描述性的術語講，它實現(xiàn)了在 Python 和 C++ 間轉換的綁定代碼（binding code）因悲；另外還有一些相當重要的 PyTorch 部分堕汞，比如 autograd 引擎和 JIT 編譯器。它也包含 C++ 前端代碼晃琳。
- aten/：這是「A Tensor Library」的縮寫（由 Zachary DeVito 命名）讯检，是一個實現(xiàn)張量運算的 C++ 庫。如果你檢查某些核代碼所處的位置卫旱，很可能就在 ATen人灼。ATen 本身就分為兩個算子區(qū)域：「原生」算子（算子的現(xiàn)代的 C++ 實現(xiàn)）和「傳統(tǒng)」算子（TH、THC顾翼、THNN投放、THCUNN），這些是遺留的 C 實現(xiàn)适贸。傳統(tǒng)的算子是其中糟糕的部分灸芳；如果可以，請勿在上面耗費太多時間取逾。
- c10/：這是「Caffe2」和「A"Ten"」的雙關語耗绿，包含 PyTorch 的核心抽象，包括張量和存儲數(shù)據(jù)結構的實際實現(xiàn)砾隅。
找代碼需要看很多地方误阻；我們應該簡化目錄結構，就是這樣晴埂。如果你想研究算子究反，你應該在 aten 上花時間。
我們看看在實踐中是如何分離這些代碼的：
當你調(diào)用一個函數(shù)時儒洛，比如 torch.add精耐，會發(fā)生什么？如果你記得我們的有關調(diào)度的討論琅锻，你腦中應該已有了這些基礎：
- 我們必須從 Python 國度轉換到 C++ 國度（Python 參數(shù)解析）卦停。
- 我們處理變量調(diào)度（VariableType—Type向胡，順便一提，和編程語言類型并無特別關聯(lián)惊完，只是一個用于執(zhí)行調(diào)度的小工具）僵芹。
- 我們處理設備類型/布局調(diào)度（Type）。
- 我們有實際的核小槐，這要么是一個現(xiàn)代的原生函數(shù)拇派，要么是傳統(tǒng)的 TH 函數(shù)。
其中每一步都具體對應于一些代碼凿跳。讓我們開路穿過這片叢林件豌。
我們在 C++ 代碼中的起始著陸點是一個 Python 函數(shù)的 C 實現(xiàn)，我們已經(jīng)在 Python 那邊見過它控嗜，像是 torch._C.VariableFunctions.add茧彤。THPVariable_add 就是這樣一個實現(xiàn)。
對于這些代碼躬审，有一點很重要：這些代碼是自動生成的棘街。如果你在 GitHub 庫中搜索，你沒法找到它們承边，因為你必須實際 build PyTorch 才能看到它們遭殉。另外一點也很重要：你不需要真正深入理解這些代碼是在做什么，你應該快速瀏覽它博助，知道它的功能险污。
我在上面用藍色標注了最重要的部分：你可以看到這里使用了一個 PythonArgParser 類來從 Python args 和 kwargs 取出 C++ 對象；然后我們調(diào)用一個 dispatch_add 函數(shù)（紅色內(nèi)聯(lián)）富岳；這會釋放全局解釋器鎖蛔糯，然后調(diào)用在 C++ 張量自身上的一個普通的舊方法。在其回來的路上窖式，我們將返回的 Tensor 重新包裝進 PyObject蚁飒。
（這里幻燈片中有個錯誤：我應該講解變量調(diào)度代碼。我這里還沒有修復萝喘。某些神奇的事發(fā)生了淮逻，于是……）
當我們在 Tensor 類上調(diào)用 add 方法時，還沒有虛調(diào)度發(fā)生阁簸。相反爬早，我有一個內(nèi)聯(lián)方法，其調(diào)用了一個內(nèi)聯(lián)方法启妹，其會在「Type」對象上調(diào)用一個虛方法筛严。這個方法是真正的虛方法（這就是我說 Type 只是一個讓你實現(xiàn)動態(tài)調(diào)度的「小工具」的原因）。
在這個特定案例中饶米，這個虛調(diào)用會調(diào)度到在一個名為 TypeDefault 的類上的 add 的實現(xiàn)桨啃。這剛好是因為我們有一個對所有設備類型（CPU 和 CUDA）都一樣的 add 的實現(xiàn)车胡；如果我們剛好有不同的實現(xiàn)，我們可能最終會得到 CPUFloatType::add 這樣的結果照瘾。正是這種虛方法的實現(xiàn)能讓我們最終得到實際的核代碼吨拍。
也希望這張幻燈片很快過時；Roy Li 正在研究使用另一種機制替代 Type 調(diào)度网杆，這能讓我們更好地在移動端上支持 PyTorch。
值得再次強調(diào)伊滋，一直到我們到達核碳却，所有這些代碼都是自動生成的。
道路蜿蜒曲折笑旺，一旦你能基本上把握方向了昼浦，我建議你直接跳到核部分。
編寫核（kernel）
PyTorch 為有望編寫核的人提供了大量有用工具筒主。在這一節(jié)我們會了解其中一些关噪。但首先，編寫核需要什么乌妙？
我們一般將 PyTorch 中的核看作由以下部分組成：
- 首先有一些我們要寫的有關核的元數(shù)據(jù)使兔，這能助力代碼生成并讓你獲取所有與 Python 的捆綁包，同時無需寫任何一行代碼藤韵。
- 一旦你到達了核虐沥，你就經(jīng)過了設備類型/布局調(diào)度。你首先需要寫的是錯誤檢查泽艘，以確保輸入的張量有正確的維度欲险。（錯誤檢查真正很重要！不要吝惜它Ｆヤ獭）
- 接下來天试，我們一般必須分配我們將要寫入輸出的結果張量。
- 該到寫核的時候了∪坏停現(xiàn)在你應該做第二次 dtype 調(diào)度喜每，以跳至其所操作的每個 dtype 特定的核。（你不應該過早做這件事脚翘，因為那樣的話你就會毫無用處地復制在任何情況下看起來都一樣的代碼灼卢。）
- 大多數(shù)高性能核都需要某種形式的并行化，這樣就能利用多 CPU 系統(tǒng)了来农。（CUDA 核是「隱式」并行化的鞋真，因為它們的編程模型構建于大規(guī)模并行化之上。）
- 最后沃于，你需要讀取數(shù)據(jù)并執(zhí)行你想做的計算涩咖！
在后面的幻燈片中海诲，我將介紹 PyTorch 中能幫你實現(xiàn)這些步驟的工具。
要充分利用 PyTorch 的代碼生成能力檩互，你需要為你的算子寫一個模式（schema）特幔。這個模式能提供你的函數(shù)的 mypy 風格類型，并控制是否為 Tensor 上的方法或函數(shù)生成捆綁包闸昨。你還可以告訴模式針對給定的設備-布局組合蚯斯，應該調(diào)用你的算子的哪種實現(xiàn)。
有關這種格式的更多信息饵较，請參閱：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/README.md
你可能也需要為你在 derivatives.yaml 中的操作定義一個導數(shù)拍嵌。
錯誤檢查可以在低層 API 完成，也能通過高層 API 實現(xiàn)循诉。低層 API 只是一個宏 TORCH_CHECK横辆，其接收的是一個布爾值，然后還有任意數(shù)量的參數(shù)構成錯誤字符串（error string）以便得出結論看該布爾值是否為真茄猫。
這個宏有個很好的地方：你可以將字符串與非字符串數(shù)據(jù)混合起來狈蚤；每一項都使用它們的 operator<< 實現(xiàn)進行格式化，PyTorch 中大多數(shù)重要的數(shù)據(jù)類型都有 operator<< 實現(xiàn)划纽。
高層 API 能讓你免于反復編寫重復的錯誤消息脆侮。其工作方法是；你首先將每個張量包裝為 TensorArg阿浓，這包含有關張量來處的信息（比如其參數(shù)名稱）他嚷。然后它提供了一些預先裝好的用于檢查多種屬性的函數(shù)；比如 checkDim() 測試的是張量的維度是否是一個固定數(shù)值芭毙。如果不是筋蓖，該函數(shù)就基于 TensorArg 元數(shù)據(jù)提供一個用戶友好的錯誤消息。
在用 PyTorch 寫算子時退敦，有一點很重要：你往往要注冊三個算子：abs_out（其操作的是一個預分配的輸出粘咖，其實現(xiàn)了 out= keyword 參數(shù)）、abs_（其操作的是 inplace）侈百、abs（這只是一個算子的普通的舊功能版本）瓮下。
大部分時間，abs_out 是真正的主力钝域，abs 和 abs_ 只是圍繞 abs_out 的薄弱 wrapper讽坏；但有時候也可為每個案例編寫專門的實現(xiàn)。
要執(zhí)行 dtype 調(diào)度例证，你應該使用 AT_DISPATCH_ALL_TYPES 宏路呜。這會獲取你想要進行調(diào)度操作的張量的 dtype，并還會為可從該宏調(diào)度的每個 dtype 指定一個 lambda。通常而言胀葱，這個 lambda 只是調(diào)用一個模板輔助函數(shù)漠秋。
這個宏不只是「執(zhí)行調(diào)度」，它也會決定你的核將支持的 dtype抵屿。這樣庆锦，這個宏實際上就有相當多一些版本，這能讓你選取不同的 dtype 子集以生成特定結果轧葛。大多數(shù)時候搂抒，你只需要 AT_DISPATCH_ALL_TYPES，但也要關注你可能需要調(diào)度其它更多類型的情況尿扯。
在 CPU 上燕耿，你通常需要并行化你的代碼。過去姜胖，這通常是通過直接在你的代碼中添加 OpenMP pragma 來實現(xiàn)。
某些時候淀散，你必須真正訪問數(shù)據(jù)右莱。PyTorch 為此提供了相當多一些選擇。
- 如果你只想獲取某個特定位置的值档插，你應該使用 TensorAccessor慢蜓。張量存取器就像是一個張量，但它將張量的維度和 dtype 硬編碼為了模板參數(shù)郭膛。當你檢索一個存取器時晨抡，比如 x.accessor
();，我們會做一次運行時間測試以確保張量確實是這種格式则剃；但那之后耘柱，每次存取都不會被檢查。張量存取器能正確地處理步幅棍现，因此你最好使用它們调煎，而不是原始的指針訪問（不幸的是，很多傳統(tǒng)的核是這樣做的）己肮。另外還有 PackedTensorAccessor士袄，這特別適用于通過 CUDA launch 發(fā)送存取器，這樣你就能從你的 CUDA 核內(nèi)部獲取存取器谎僻。（一個值得一提的問題：TensorAccessor 默認是 64 位索引娄柳，這比 CUDA 中的 32 位索引要慢得多！） - 如果你在用很常規(guī)的元素存取編寫某種算子艘绍，比如逐點運算赤拒，那么使用遠遠更高級的抽象要好得多，比如 TensorIterator。這個輔助類能為你自動處理廣播和類型提升（type promotion）需了，相當好用跳昼。
- 要在 CPU 上獲得真正的速度，你可能需要使用向量化的 CPU 指令編寫你的核肋乍。我們也有用于這方面的輔助函數(shù)鹅颊！Vec256 類表示一種標量向量，并提供了一些能在它們上一次性執(zhí)行向量化運算的方法墓造。然后 binary_kernel_vec 等輔助函數(shù)能讓你輕松地運行向量化運算堪伍，然后結束那些沒法用普通的舊指令很好地轉換成向量指令的東西。這里的基礎設施還能在不同指令集下多次編譯你的核觅闽，然后在運行時間測試你的 CPU 支持什么指令帝雇，再在這些情況中使用最佳的核。
PyTorch 中大量核都仍然是用傳統(tǒng)的 TH 風格編寫的蛉拙。（順便一提尸闸，TH 代表 TorcH。這是個很好的縮寫詞孕锄，但很不幸被污染了吮廉；如果你看到名稱中有 TH，可認為它是傳統(tǒng)的畸肆。）傳統(tǒng) TH 風格是什么意思呢宦芦？
- 它是以 C 風格書寫的，沒有（或很少）使用 C++轴脐。
- 其 refcounted 是人工的（使用了對 THTensor_free 的人工調(diào)用以降低你使用張量結束時的 refcounts）调卑。
- 其位于 generic/ 目錄，這意味著我們實際上要編譯這個文件很多次大咱，但要使用不同的 #define scalar_t
這種代碼相當瘋狂恬涧，而且我們討厭回顧它，所以請不要添加它碴巾。如果你想寫代碼但對核編寫了解不多气破，你能做的一件有用的事情：將某些 TH 函數(shù)移植到 ATen。
工作流程效率
最后我想談談在 PyTorch 上的工作效率餐抢。如果 PyTorch 那龐大的 C++ 代碼庫是阻攔人們?yōu)?PyTorch 做貢獻的第一只攔路虎现使，那么你的工作流程的效率就是第二只。如果你想用 Python 習慣開發(fā) C++旷痕，那可能會很艱辛：重新編譯 PyTorch 需要大量時間碳锈，你也需要大量時間才能知道你的修改是否有效。
如何高效工作本身可能就值得做一場演講欺抗，但這頁幻燈片總結了一些我曾見過某些人抱怨的最常見的反模式：「開發(fā) PyTorch 很困難售碳。」
- 如果你編輯一個 header，尤其是被許多源文件包含的 header（尤其當被 CUDA 文件包含時）贸人，可以預見會有很長的重新 build 時間间景。盡量只編輯 cpp 文件，編輯 header 要審慎艺智！
- 我們的 CI 是一種非常好的零設置的測試修改是否有效的方法倘要。但在獲得返回信號之前你可能需要等上一兩個小時。如果你在進行一種將需要大量實驗的改變十拣，那就花點時間設置一個本地開發(fā)環(huán)境封拧。類似地，如果你在特定的 CI 配置上遇到了困難的 debug 問題夭问，就在本地設置它泽西。你可以將 Docker 鏡像下載到本地并運行：https://github.com/pytorch/ossci-job-dsl
- 貢獻指南解釋了如何設置 ccache：https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#use-ccache ；強烈建議這個缰趋，因為這可以讓你在編輯 header 時幸運地避免大量重新編譯捧杉。當我們在不應該重新編譯文件時重新編譯時，這也能幫你覆蓋我們的 build 系統(tǒng)的漏洞秘血。
- 最后糠溜，我們會有大量 C++ 代碼。如果你是在一臺有 CPU 和 RAM 的強大服務器上 build直撤，那么會有很愉快的體驗。特別要說明蜕着，我不建議在筆記本電腦上執(zhí)行 CUDA build谋竖。build CUDA 非常非常慢，而筆記本電腦往往性能不足承匣，不足以快速完成蓖乘。
參與進來！
這就是我們旋風一般的 PyTorch 內(nèi)核之旅了韧骗！其中省略了很多很多東西嘉抒；但希望這里的描述和解釋至少能幫你消化其代碼庫中相當大一部分。
接下來該做什么袍暴？你能做出怎樣的貢獻些侍？我們的問題跟蹤器是個開始的好地方：https://github.com/pytorch/pytorch/issues。
從今年開始政模，我們一直在分類鑒別問題岗宣；標注有「triaged」的問題表示至少有一個 PyTorch 開發(fā)者研究過它并對該問題進行了初步評估。你可以使用這些標簽找到我們認為哪些問題是高優(yōu)先級的或查看針對特定模塊（如 autograd）的問題淋样，也能找到我們認為是小問題的問題耗式。（警告：我們有時是錯的！）
即使你并不想馬上就開始寫代碼，也仍有很多其它有用的工作值得去做刊咳，比如改善文檔（我很喜歡合并文檔 PR彪见，它們都很贊）、幫助我們重現(xiàn)來自其他用戶的 bug 報告以及幫助我們討論問題跟蹤器上的 RFC娱挨。沒有我們的開源貢獻者余指，PyTorch 不會走到今天；我們希望你也能加入我們让蕾！
