API統一、乾淨,新型EagerPy實現多框架無縫銜接

選自arXiv

編輯：杜偉、小舟

這個新型 Python 框架對庫開發者和用戶都大有裨益。

近年來，深度學習領域的進展與深度學習框架的開發同步進行。這些框架為自動微分和 GPU 加速提供了高級且高效的 API，從而可以利用相對較少和簡單的代碼實現極度複雜和強大的深度學習模型。

最初，Theano、Caffe、MXNet、TensorFlow 和 CNTK 等很多流行的深度學習框架使用的是基於圖的方法。用戶首先需要定義一個靜態數據流圖（static data flow graph），然後可以對它進行高效地微分、編譯並在 GPU 上執行。所以，提前了解整個計算圖有助於實現高性能。

但是，這種方法導致難以調試模型以及實現具有變化圖（changing graph）的動態模型（如 RNN）。

所以，針對這種方法的局限性，深度學習模型的 Eager Execution 成為了深度學習研究領域的主流方法。用戶不再需要提前構建靜態數據流圖，Eager Execution 框架自身就可以提供 define-by-run 的 API，它可以高速地構建臨時的動態圖。目前，兩大主流深度學習框架 PyTorch 和 TensorFlow 都在使用 eager execution 方法。

而在本文中，來自德國圖賓根大學和圖賓根伯恩斯坦計算神經科學中心的研究者將 eager execution 進行了擴展，提供了一個新的 Python 框架 EagerPy，它可以編寫自動且原生地適配 PyTorch、TensorFlow、Jax 和 Numpy 的代碼。EagerPy 對庫開發者和用戶都有裨益。

論文地址：https://arxiv.org/abs/2008.04175v1

項目地址：https://github.com/jonasrauber/eagerpy

EagerPy 能夠編寫與框架無關（framework-agnostic）的代碼，這些代碼可以與 PyTorch、TensorFlow、Jax 和 NumPy 實現原生地適配。

這樣一來，首先對於新庫開發者而言，他們不僅可以選擇同時支持上述這幾個主流深度學習框架或者為每個框架重新實現庫，而且可以對代碼重複進行處理。

其次對於這些庫的使用者而言，他們也可以更輕鬆地切換深度學習框架，並且不會被特定的第三方庫鎖定。

不僅如此，單個框架的使用者也會從 EagerPy 中獲益，這是因為 EagerPy 提供了全面的類型注釋以及對方法連結到任何框架的一致支持。

接下來我們來看 EagerPy 的具體設計與實現。

EagerPy 的設計與實現

EagerPy 的構建考慮到了 4 個設計目標。兩個主要的目標是為需要執行操作的人提供統一的 API，並維護底層框架的原始性能。這兩個主要目標定義了 EagerPy 是什麼，所以是設計的核心。

與底層框架特定的 API 相比，完全可連結的 API 和全面的類型檢查支持這兩個附加目標使 EagerPy 更加易於使用，也更安全。

儘管進行了這些更改和改進，但研究者嘗試避免不必要的熟悉度（familiarity）損失。只要有意義，EagerPy API 都會遵循 NumPy、PyTorch 和 JAX 設置的標準。

統一的 API

為了實現語法上的一致性，研究者使用適當的方法定義了一個抽象 Tensor 類，並使用一個實例變量來保存原生張量（native tensor），然後為每個支持的框架實現一個特定的子類。對於諸如 sum 或 log 的很多操作，這就像調用底層框架一樣簡單；而對於其他操作，則工作量會稍大一些。

最困難的部分是統一自動微分 API。PyTorch 使用了一個低級的 autograd API，該 API 允許但也需要對反向傳播的精確控制。TensorFlow 使用基於梯度磁帶（gradient tapes）的更高級 API。而 JAX 使用基於微分函數的相當高級的 API。

所以，為了統一它們，EagerPy 模仿了 JAX 的高級功能 API，並在 PyTorch 和 TensorFlow 中重新實現。EagerPy 通過 value_and_grad_fn() 函數將其開放。

此外，能夠編寫自動與所有支持的框架一起運行的代碼，不僅需要語法，還需要語義統一。為了確保這一點，EagerPy 附帶了一個龐大的測試套件，該套件可以驗證不同框架特定子類之間的一致性。它會在所有 pull-request 上自動運行，並且需要通過之後才能合併新代碼。

測試套件還可以作為所支持的操作和參數組合的最終參考。這樣就可以避免文檔和實現之間出現不一致，並在實踐中引出測試驅動開發過程。

原始性能

沒有 EagerPy，想要與不同深度學習框架進行交互的代碼必須經過 NumPy 實現。這需要在 CPU（NumPy）和 GPU（PyTorch、TensorFlow 和 JAX）之間進行高成本的內存複製，反之亦然。

此外，許多計算僅在 CPU 上執行，為了避免這種情況，EagerPy 僅保留對原始框架特定張量的引用（例如 GPU 上的 PyTorch 張量），並將所有的操作委託給相應的框架。這幾乎不產生任何的計算開銷。

完全可連結的 API

求和或平方之類的許多運算都要採用張量並返回一個張量。通常情況下，這些運算按順序被調用。例如使用平方、求和和開平方根以計算 L2 範數。

在 EagerPy 中，所有運算都成為了張量對象（tensor object）上可用的方法。這樣就可以按照它們的自然順序（x.square().sum().sqrt()）來連結操作。相反，例如，NumPy 需要相反的操作順序，即 np.sqrt(np.square(x).sum())。

類型檢查

在 Python3.5 中，Python 語法的擴展已經實現了對類型注釋的支持（van Rossum 等人，2015 年）。即使具有類型注釋，Python 仍然是一種動態類型化的程式語言，並且當前在運行時會忽略所有類型注釋。但是，我們可以在運行代碼之前通過靜態代碼分析器檢查這些類型注釋。

EagerPy 帶有所有參數和返回值的全面類型注釋，並使用 Mypy（Lehtosalo 等人，2016 年）對這些注釋進行檢查。這有助於我們捕獲 EagerPy 中的漏洞，否則這些漏洞將一直不會被發現。

EagerPy 用戶可以通過鍵入自己代碼的注釋，並根據 EagerPy 的函數籤名（function signature）自動檢查代碼來進一步優化。這一點很關鍵，因為 TensorFlow、NumPy 和 JAX 當前自身不提供類型注釋。

EagerPy 的代碼實例解析

如下代碼 1 為一個通用 EagerPy 範數函數，它可以通過任何框架中的原生張量被調用，並且返回的範數依然作為同一個框架中的原生張量。

代碼 1：框架無關的範數函數。

EagerPy 和原生張量之間的轉換

原生張量可以是 PyTorch GPU 或 CPU 張量，如下代碼 2 所示：

代碼 2：原生 PyTorch 張量。

可以是 TensorFlow 張量，如下代碼 3 所示：

代碼 3：原生 TensorFlow 張量。

可以是 JAX 數組，如下代碼 4 所示：

代碼 4：原生 JAX 數組。

可以是 NumPy 數組，如下代碼 5 所示：

代碼 5：原生 NumPy 數組。

無論是哪種原生張量，通常都可以使用 ep.astensor 將它轉換為適當的 EagerPy 張量。在此步驟中，通過使用正確的 EagerPy 張量類來自動封裝原生張量。此外，最初的原生張量通常可以利用. raw 屬性實現訪問。完整示例如下代碼 6 所示：

EagerPy 和原生張量之間的轉換。

在函數中通常將所有輸入轉換為 EagerPy 張量。這可以通過單獨調用 ep.astensor 完成，但在使用 ep.astensors 時，代碼可以更加簡潔，如下：

實現框架無關的通用函數

通過上文中的轉換函數，我們可以定義一個簡單的框架無關函數，如下代碼 8 所示：

代碼 8：一個簡單的框架無關範數函數。

如下代碼 9 所示，通過一個 PyTorch 張量來調用範數函數：

如下代碼 10 所示，通過一個 TensorFlow 張量來調用範數函數：

此外，還需要注意一點，如果如上代碼 8 所示使用 EagerPy 張量來調用函數，則 ep.astensor 調用只會返回它的輸入。但是，最後一行代碼中的 result.raw 調用依然會提取底層原生張量。通常而言，實現的通用函數最好可以透明地操控任何原生張量和 EagerPy 張量，也就是說返回類型應該總是與輸入類型相匹配。

這在 Foolbox 等庫中非常有用，可以使用戶同時處理 EagerPy 和原生張量。

為此，EagerPy 提供上述轉換函數的兩種派生函數，分別是 ep.astensor_和 ep.astensors_，它們可以返回一個能夠恢復輸入類型的反轉函數。

如果 astensor_的輸入是一個原生張量，則 restore_type 等同於. raw；而如果原輸入是一個 EagerPy 張量，則 restore_type 將不會調用. raw。因此，我們可以編寫對任何輸入都透明的改進版框架無關通用函數，如下代碼 11 所示：

最後，如下代碼 12 所示，使用 ep.astensors_來轉換和恢復多個輸入：