張量形狀不匹配是深度神經(jīng)網(wǎng)絡(luò)機(jī)器學(xué)習(xí)過程中會出現(xiàn)的重要錯(cuò)誤之一。由于神經(jīng)網(wǎng)絡(luò)訓(xùn)練成本較高且耗時(shí)，在執(zhí)行代碼之前運(yùn)行靜態(tài)分析，要比執(zhí)行然后發(fā)現(xiàn)錯(cuò)誤快上很多。

由于靜態(tài)分析是在不運(yùn)行代碼的前提下進(jìn)行的，因此可以幫助軟件開發(fā)人員、質(zhì)量保證人員查找代碼中存在的結(jié)構(gòu)性錯(cuò)誤、安全漏洞等問題，從而保證軟件的整體質(zhì)量。

相比于程序動態(tài)分析，靜態(tài)分析具有不實(shí)際執(zhí)行程序；執(zhí)行速度快、效率高等特點(diǎn)而廣受研究者青睞，目前，已有許多分析工具可供研究使用，如斯坦福大學(xué)開發(fā)的 Meta-Compilation（Coverity）、利物浦大學(xué)開發(fā)的 LDRA Testbed 等。

近日，來自韓國首爾大學(xué)的研究者們提出了另一種靜態(tài)分析器 PyTea，它可以自動檢測 PyTorch 項(xiàng)目中的張量形狀錯(cuò)誤。在對包括 PyTorch 存儲庫中的項(xiàng)目以及 StackOverflow 中存在的張量錯(cuò)誤代碼進(jìn)行測試。結(jié)果表明，PyTea 可以成功的檢測到這些代碼中的張量形狀錯(cuò)誤，幾秒鐘就能完成。

• 論文地址：https://arxiv.org/pdf/2112.09037.pdf
• 項(xiàng)目地址：https://github.com/ropas/pytea

幾秒就能查找張量形狀錯(cuò)誤的 PyTea

PyTea 工具可以靜態(tài)地掃描 PyTorch 程序并檢測可能的形狀錯(cuò)誤。PyTea 通過額外的數(shù)據(jù)處理和一些庫（例如 Torchvision、NumPy、PIL）的混合使用來分析真實(shí)世界 Python/PyTorch 應(yīng)用程序的完整訓(xùn)練和評估路徑。

PyTea 的工作原理是這樣的：給定輸入的 PyTorch 源，PyTea 靜態(tài)跟蹤每個(gè)可能的執(zhí)行路徑，收集路徑張量操作序列所需的張量形狀約束，并決定約束滿足與否（因此可能發(fā)生形狀錯(cuò)誤）。

具體來說：如下圖所示， PyTea 首先將原始 Python 代碼翻譯成一種內(nèi)核語言，即 PyTea 內(nèi)部表示（PyTea IR）。然后，它跟蹤轉(zhuǎn)換后的 IR 的每個(gè)可能執(zhí)行路徑，并收集有關(guān)張量形狀的約束，這些約束規(guī)定了代碼在沒有形狀錯(cuò)誤的情況下運(yùn)行的條件。 PyTea 將收集到的約束集提供給 SMT（Satisfiability Modulo Theories）求解器 Z3，以判斷這些約束對于每個(gè)可能的輸入形狀都是可滿足的。根據(jù)求解器的結(jié)果，PyTea 會得出結(jié)論，哪條路徑包含形狀錯(cuò)誤。如果 Z3 的約束求解花費(fèi)太多時(shí)間，PyTea 會停止并發(fā)出「don’t know」提示。

PyTea 的整體結(jié)構(gòu)。

PyTea 由兩個(gè)分析器組成，在線分析器：node.js (TypeScript / JavaScript)；離線分析器：Z3 / Python。

• 在線分析器：查找基于數(shù)值范圍的形狀不匹配和 API 參數(shù)的濫用。如果 PyTea 在分析代碼時(shí)發(fā)現(xiàn)任何錯(cuò)誤，它將停在該位置并將錯(cuò)誤和違反約束通知用戶；
• 離線分析器：生成的約束傳遞給 Z3 。Z3 將求解每個(gè)路徑的約束集并打印第一個(gè)違反的約束（如果存在）。

我們先來看下結(jié)果展示，在線分析器發(fā)現(xiàn)錯(cuò)誤：

離線分析器發(fā)現(xiàn)錯(cuò)誤：

為了更好的理解 PyTea 執(zhí)行靜態(tài)分析過程，下面我們介紹一下主要的技術(shù)細(xì)節(jié)，包括 PyTorch 程序結(jié)構(gòu)、張量形狀錯(cuò)誤、PyTea IR 等，以便讀者更好的理解執(zhí)行過程。

首先是 PyTorch 程序結(jié)構(gòu)，PyTorch、TensorFlow 和 Keras 等現(xiàn)代機(jī)器學(xué)習(xí)框架需要使用 Python API 來構(gòu)建神經(jīng)網(wǎng)絡(luò)。使用此類框架訓(xùn)練神經(jīng)網(wǎng)絡(luò)大多遵循如下四個(gè)階段的標(biāo)準(zhǔn)程序。

在 PyTorch 中，常規(guī)神經(jīng)網(wǎng)絡(luò)訓(xùn)練代碼的結(jié)構(gòu)。

訓(xùn)練模型需要先定義網(wǎng)絡(luò)結(jié)構(gòu)，圖 2 為一個(gè)簡化的圖像分類代碼，取自官方的 PyTorch MNIST 分類示例：

在這里，上述代碼首先定義一系列神經(jīng)網(wǎng)絡(luò)層，并使它們成為單一的神經(jīng)網(wǎng)絡(luò)模塊。為了正確組裝層，前一層的返回張量必須滿足下一層的輸入要求。網(wǎng)絡(luò)使用超參數(shù)的初始化參數(shù)進(jìn)行實(shí)例化，例如隱藏層的數(shù)量。接下來，對輸入數(shù)據(jù)集進(jìn)行預(yù)處理并根據(jù)網(wǎng)絡(luò)的要求進(jìn)行調(diào)整。從該階段開始，每個(gè)數(shù)據(jù)集都被切成較小的相同大小的塊（minibatch）。最后，主循環(huán)開始，minibatch 按順序輸入網(wǎng)絡(luò)。一個(gè) epoch 是指將整個(gè)數(shù)據(jù)集傳遞到網(wǎng)絡(luò)的單個(gè)循環(huán)，并且 epoch 的數(shù)量通常取決于神經(jīng)網(wǎng)絡(luò)的目的和結(jié)構(gòu)。除了取決于數(shù)據(jù)集大小的主訓(xùn)練循環(huán)之外，包括 epoch 數(shù)在內(nèi)，訓(xùn)練代碼中的迭代次數(shù)在大多數(shù)情況下被確定為常數(shù)。

在構(gòu)建模型時(shí)，網(wǎng)絡(luò)層之間輸入、輸出張量形狀的不對應(yīng)就是張量形狀錯(cuò)誤。通常形狀錯(cuò)誤很難手動查找，只能通過使用實(shí)際輸入運(yùn)行程序來檢測。下圖就是典型的張量形狀錯(cuò)誤（對圖 2 的簡單修改），如果不仔細(xì)查看，你根本發(fā)現(xiàn)不了錯(cuò)誤：

對于張量形狀錯(cuò)誤（如上圖的錯(cuò)誤類型），PyTea 將原始 Python 代碼翻譯成 PyTea IR 進(jìn)行查找，如下圖是 PyTea IR 示例：

上面提到，PyTea 會跟蹤轉(zhuǎn)換后的 IR 的每個(gè)可能執(zhí)行路徑，并收集有關(guān)張量形狀約束。其實(shí)約束是 PyTorch 應(yīng)用程序所需要的條件，以便在沒有任何張量形狀誤差的情況下執(zhí)行它。例如，一個(gè)矩陣乘法運(yùn)算的兩個(gè)操作數(shù)必須共享相同的維數(shù)。下圖顯示了約束的抽象語法：

約束的抽象語法部分截圖

如何使用 PyTea

首先，安裝環(huán)境要求：node.js >= 12.x，python >= 3.8，z3-solver >= 4.8。

安裝和使用可參考以下代碼：

# install node.js 
sudo apt-get install nodejs 
# install python z3-solver 
pip install z3-solver 
# download pytea 
wget https://github.com/ropas/pytea/releases/download/v0.1.0/pytea.zip 
unzip pytea.zip 
# run pytea 
python bin/pytea.py path/to/source.py 
# run example file 
python bin/pytea.py packages/pytea/pytest/basics/scratch.py 

編譯代碼：

# install dependencies 
npm run install:all 
pip install z3-solver 
# build 
npm run build 

網(wǎng)友紛紛表示：這個(gè)工具絕對是必要的！