treevalue——Master Nested Data Like Tensor

HansBug發表於2022-02-14

首先,請和我一起高呼——“treevalue——通用樹形結構建模工具 + 極簡樹形結構程式設計模型”。

咳咳,好久沒更新了,這一次是真的好久不見,甚是想念。在之前的三期中,關於 treevalue 的核心特性等內容已經基本完成了講述。因此本篇作為該系列的終章,將嘗試用更高一層的視角來分析 treevalue ,以求進一步瞭解其核心思想與應用模式,並通過乾貨資料與例項展示來展現其真實能力。

閒話少敘,讓我們開始吧!如果還沒了解過 treevalue 的小夥伴們可以先去讀一下之前的幾篇文章:

整體設計

概念與架構

想要完整地瞭解 treevalue ,首先還是需要從整體架構上來看看,如下圖所示。

(treevalue的整體架構)

顯而易見的一點, treevalue 是基於Python構建的,並且處於運算效能上的考慮,因此也大量採用了Cython來實現,該工具可以用一種介於Python和C/C++之間的語法進行編碼,並以C/C++的形式編譯為靜態庫,通過繞開一系列不必要的動態機制等方式來實現明顯的加速。

treevalue 中,最底層的為 TreeStorage ,為資料層,主要對樹狀資料結構進行管理,並對上層提供最基本的介面。建立在之上的 TreeValue 為最關鍵的一個類,基於資料層進行了基本的封裝,能實現 treevalue 的基本特性且具備進一步擴充套件的能力。樹化(treelize)在之前的幾篇文章中有過介紹,其作用在於將現有庫的函式和類進行樹化的擴充套件,使 treevalue 的特性可以被快速應用至現有工具上。工具部分(utilities)即為基於 TreeValue 構建的一些簡單工具,支援了基本的樹操作、函式式運算等功能。另一個極為重要的部分即為 FastTreeValue 類,其為 TreeValue 的子類,包含了大部分的運算子,並可以實現“裝載即用”,只需要將現有的物件裝入FastTreeValue即可批量訪問其屬性、批量呼叫其方法,因此 FastTreeValue 被最為廣泛的使用。這裡需要注意的是, treevalue 並未針對任何一個特定的現有庫進行特殊化設計,而是通過泛用型設計,讓諸如PyTorch、Numpy、Tensorflow在內的幾乎全部介面均可以被快速擴充套件到樹運算上

總的來說,基於 treevalue ,可以將現有輪子中的函式與類擴充套件為支援樹狀運算的新形態,並可以讓程式設計者基於這一擴充套件進行更加方便快捷的程式碼構建。

設計定位

在前面幾篇該系列的文章裡,我們都加上了“強化學習(Reinforcement Learning)”的標籤,也是因此,此文才會被推送到關注了這一話題的使用者手裡。可是你們大概已經發現,無論是之前的三篇文章,還是我們的原始碼倉庫,都並沒有深度強化學習相關的內容,甚至在原始碼中看不到 numpytorch 等常見AI相關庫的存在。想必各位讀者也從之前的文章發出後便早已疑惑多時了。因此在這裡我們來回答一個重要的問題—— treevalue的設計定位是什麼

其實仔細琢磨的話,答案也已經很明顯了—— treevalue從設計之初就是一款具備充分泛用型,且易用性為主執行效能為輔的樹資料結構計算框架。之前各個文章乃至程式碼中未大量涉及具體的深度強化學習內容也正是這個原因。在Treevalue(0x02)——函式樹化詳細解析(上篇)中介紹的 func_treelize 特性也是可以針對任意函式進行使用的,因此即便不是 torch 這樣的經典的庫,換上其他的庫與其他的類、方法,也同樣可以使用treevalue,且使用體驗與原本的庫基本一致,極為友好。甚至於深度強化學習以外的任意領域,只要有基於樹結構的運算的地方, treevalue 就有它的用武之地——這也正是它設計上的核心理念,與真正的優勢所在。

與同類產品的對比

基於上述的介紹,讀者們可能會擔心一個易用性、泛用性為主的庫,會不會在其他上,尤其是在深度強化學習程式的特性支援與執行效能上存在劣勢。這樣的顧慮其實很正常,因為在大部分情況下,軟體的泛用性和效能都是不容易做到兼顧的。但是 treevalue 是否存在這樣的問題呢?因此本節將會通過與同類或相近產品的對比分析,來對這一問題進行回答

同類或相近產品介紹

treevalue 開發與測試的階段,我們發現其實在現有的其他庫中,也有類似的樹狀運算支援。其中比較具有典型意義的有以下四個:

  • dm-tree,DeepMind團隊開發的輕量級樹狀運算庫。
  • Tianshou-Batch,清華大學機器學習組開發,基於PyTorch的深度強化學習庫,其中Batch為重點設計的一個關鍵元件。
  • jax-libtree,谷歌團隊開發的機器學習庫,其中libtree為內部的一個輕量級樹運算封裝庫。
  • torchbeast-nest,Meta研究院(原facebook研究院)開發,基於PyTorch的深度強化學習框架,其中nest為內部的一個輕量級樹運算封裝庫。

特性對比

基於上面提到的四款同類產品及 treevalue ,所開源的程式碼及提供的文件,可以得出以下的分析。

同類產品 重量級 泛用性 函式式 結構運算 自建結構 函式擴充套件 自擴充套件
dm-tree 輕量級 × × × ×
Tianshou-Batch 中輕量級 × × × × ×
jax-libtree 輕量級 × × ×
torchbeast-nest 輕量級 × × ×
treevalue 中量級

具體來說:

  • dm-tree為基於C++實現的獨立輕量級庫,不針對特定的框架,且支援簡單的函式式運算。
  • tianshou-Batch為中輕量級庫,自建了一套針對PyTorch的樹資料結構,實現了基本的增刪查改功能,並且實現了包括stack、split在內的少量運算操作。
  • jax-libtree為Google Jax庫中的子模組,不針對特定框架,支援簡單的函式式運算,還支援了諸如transpose這樣的樹結構運算。
  • torchbeast-nest為Meta torchbeast中的樹結構子模組,不針對特定框架,支援簡單的函式式運算,此外還支援簡單的二元函式擴充套件運算能力。
    而相比之下, treevalue 為一款泛用型中量級庫,且支援主要的函式式運算。其實現原理為自建樹狀資料結構,並基於這一結構展開運算。不僅如此,還支援了比jax-libtree更加豐富的結構運算(subside與支援自動檢測結構的rise),也支援了比torchbeast-nest更靈活的函式擴充套件能力(func_treelize)。而 treevalue 最為強大的地方體現在treetensor中,只需要對部分 torch.Tensor 的方法進行特別支援後,剩下的全部方法均可在現有框架上實現,並保持和原有API一樣的使用方式。

這一點,意味著對於基於 treevalue 的開發者而言,不再需要大規模逐個進行封裝遷移,只需要針對個別較特殊的API進行特別實現,其他的可以直接批量快速對映,直接做到無一遺漏。實際上,對於相容庫的開發而言,往往實現了90%,甚至與99%的介面,都可能對實際使用帶來較多的限制,而基於treevalue開發的應用而言,則在原理上便不存在這一問題,介面的擴充套件可以一步到位。而對於使用者而言,相容庫的使用體驗將和原有庫高度一致,原有的運算性質不會改變,且進一步支援了樹狀運算,而一步到位的相容更是讓使用者不會被不全的API所限制,使用體驗極佳,甚至於可以將現有程式碼以極小規模的修改便實現從原有庫向相容庫的遷移,很容易確保正常執行,並可以進一步簡化程式碼實現

效能對比

上文中提到, treevalue 是其中唯一的一款中量級庫,這意味著從實現的角度來看,其需要為了支援高度的泛用性和易用性,而考慮更多的情況,設立更多的處理機制。因此,理論上 treevalue 會在效能上存在一定的劣勢。然而事實真的如理論所言嗎?讓我們來看看接下來的效能分析資料。

在本次效能對比中,我們將根據同類產品的特點,將實驗分為兩組:

  • 資料結構組——tianshou Batch和treevalue,比較物件為split(將單個Tensor進行拆分)、stack(將多個Tensor進行拼接)運算的效能
  • 泛用運算組——dm-tree、jax-libtree、torchbeast-nest和treevalue,比較物件為flatten(將樹結構展開為可逆的列表結構)、mapping(簡單的函式式對映運算)運算的效能
    資料結構組的效能測試結果如下所示,不難發現 treevalue 在張量運算的效能上具有明顯優勢。


(treevalue和tianshou Batch在torch.split操作上的效能對比,treevalue具有明顯的優勢)


(treevalue和tianshou Batch在torch.stack操作上的效能對比,treevalue具有明顯的優勢)

泛用運算組的效能測試結果如下圖所示,其中dm-tree為其中最輕量級的庫,卻意外地擁有最低的效能;而jax-libtree和torchbeast-nest相比之下擁有不錯的效能。但是,作為中量級庫且包含大量易用性設計的treevalue,仍然在各種規模的資料上擁有效能優勢——在小規模資料上優勢極為明顯,即便在大規模資料上依然可以較之jax-libtree保持微弱優勢,且較之torchbeast-nest保持大比例的優勢。


(treevalue、dm-tree、jax-libtree、torchbeast-nest在flatten運算上的效能對比,treevalue整體處於優勢)


(treevalue、dm-tree、jax-libtree、torchbeast-nest在mapping運算上的效能對比,treevalue整體處於優勢)

不僅如此,我們還針對其他的各類基礎操作於運算進行了效能測試,更多的測試結果資訊參見專案README.md

經過一系列的測試,結果表明效能問題並未出現,甚至於相較於同類產品仍有效能優勢。這一點對於中量級設計的庫而言是極為難得的,這意味著無論從效能還是從使用體驗來說,treevalue均具備全方位的優勢

案例對比

在上文中,我們在與同類產品的的對比中,展現了 treevalue 強大的泛用性與易用性,以及堪稱優秀的運算效能。而實際上要想更具說服力地展現這一點,還是需要一些更加具體的例子。

首先是官方文件中的兩組例子,分別用於Numpy和Scikit-Learn:

  • Apply Into Numpy,這組例子展現了在實現同樣涉及樹資料結構的情況下,使用FastTreeValue的程式碼極為明顯地短於原有實現,且簡潔清晰程度遠超以往
  • Apply Into Scikit-Learn,該組例子展現了基於原有的介面,通過極簡的裝飾擴充套件即可將PCA運算應用到整棵樹上,得出整棵樹的解,全程高度精簡易懂。
    此外,我們還準備了一個 treevalue 用於具體DRL專案的例子,這次我們選擇了在深度強化學習領域最複雜的專案——AlphaStar中的相關資料處理函式 collate_fn ,這個函式的作用是在每個訓練iteration之前對資料進行堆疊、填充和預處理。原始的程式碼充斥著大量的for迴圈和if-else分支控制,而通過使用 treevalue ,我們可以將上述所有操作完全用並行操作API重寫,具體來說,兩個版本的程式碼度量分析如下表所示:

這表明, treevalue 在複雜專案上的優勢將更加顯著,具體表現為程式碼行數縮減到1/3邏輯複雜度大幅降低可維護性明顯提高——以及由此帶來的編碼用時大幅縮減。此外,在效能測試中,treevalue版程式碼相較於原始程式碼也並無劣勢。至此, treevalue 的各項能力已經得到了充分的驗證。

展望

treevalue 的各項能力,在上文中已經進行了充分的論述。因此,本節將採用開放式問題的方式,丟擲一系列問題,也作為展望。歡迎讀者一同參與討論,以及如有更進一步的腦洞,也歡迎留言~~

問題1:你認為, treevalue 僅僅只是一個庫?一組操作工具?還是一整套運算模型?

問題2treevalue運算模型的核心是什麼?此處核心指的是,基於該核心可以直接或間接衍生出幾乎全部現有特性。

問題3:在之前的文章中,已經詳細描述了函式的樹化擴充套件(treelize)機制,效果為將普通函式擴充套件為支援樹運算的函式。除此之外,是否還有其他類似的可擴充套件物件?是否有可以基於函式樹化的上位可擴充套件物件?

問題4:如果需要對現有的資料模型類(例如torch.Tensor、np.ndarray等)進行樹化擴充套件,需要注意哪些問題?工廠類(工廠模式構建的類)呢?工具類(整合靜態工具的類,多見於Java)呢?對模組(module)的擴充套件呢?上述的各種擴充套件機制是否可以歸納為另一種統一的運算模型

問題5:treevalue相較於同類產品的效能優勢雖然普遍存在,但依然沒有做到全方位拉開差距。有哪些可能的優化點?其中哪些是可以針對具體應用情況作針對性優化的?

問題6:如何通過擴充套件讓treevalue支援任意型別的資料容器(包括list、tuple,甚至自建資料模型)?眼下的主要障礙在哪裡?是在技術上還是運算模型上

問題7:你還能想到其他的 treevalue + 機器學習的應用嗎?圖神經網路資料結構的表示是否可以實現?

相關文章