機器之心原創
機器之心編輯部
現在都 2021 年了，機器學習好填的坑都已經填了，大家都在想怎麼將模型用到各種實際任務上。我們再去討論深度學習框架，吐槽它們的體驗，會不會有點過時？並不會，新模型與新算法，總是框架的第一生產力。
從 Theano 一代元老，到 TensorFlow 與 PyTorch 的兩元世界，到現在各個國產框架與工具組件的興起。深度學習框架，總是跟隨前沿 DL 技術的進步而改變。
不過今天並不是討論深度學習框架的演變，而只是單純分享一下在算法工程中，使用 TensorFlow 遇到的各種問題與感想。
TF 1.X ：令人又愛又恨
TensorFlow 2.X 已經正式發布 1 年多了，一周多前 TF 2.4 剛剛發布，看 Release Notes 最新版仍然關注多機並行訓練、Keras 性能等新模塊，甚至發布了「TensorFlow 版」的 NumPy 工具。然而，除去這些新特性，TF 2.X 很多不和諧的問題仍然存在。
以至於，一直維護 TF 1.15 的算法工程師，似乎 TensorFlow 的更新，對自己沒有任何影響。TF 1.X，仍然活躍在眾多的 GPU 上。
如果我們開始一項新任務，最先要做的就是查找已有的研究，以及已有的開原始碼。這樣的開原始碼，即使到現在，很多最新的前沿模型，尤其是谷歌大腦的各項研究，仍然採用的 1.X 的寫法與 API。
比如說，預訓練語言模型 T5、Albert、Electra 或者圖像處理模型 EfficientNet 等等。他們實際上還是用 1.X 那一套方法寫的，只不過能兼容 TensorFlow 2.X。
你會驚奇地發現，它們的 TensorFlow 導入都是這種風格：
import tensorflow.compat.v1 as tfimport tensorflow.compat.v2 as tf
其中，「compat」是 TF2.X 專門為兼容 TF 1.X 配置的模塊。目前，還是有很多前沿研究，放不下 TF 1.X。那就更不用說之前的經典模型，絕大多都是 TF 1.X 寫的。
不過如果只是導入「compat」模塊，那麼使用 TensorFlow 2.0 是為了什麼？難道只是饞它的版本號麼。
維護 OR 更新？
假設我們要使用這些 TF 模型，從開原始碼開始進行修改或重寫。那麼就遇到了第一個問題，我到底是維護一個 TF 1.X 的代碼庫呢，還是忍痛更新的 2.X？
假定我們決定維護 1.X 的靜態計算圖，那麼你會發現，我們寫代碼只是在寫計算圖，中間變量列印不出信息，循環語句或條件語句，基本都要改成矩陣運算形式。
TF 1.X 還是挺費勁的，就說列印變量，只調用 tf.print() 還不行，你還有將這條語句綁定到主要計算流程上，控制 Dependency。這樣，才能列印出真實的矩陣信息。
假設我們選擇更新到 TF 2.0，基本上就相當於重寫模型了。官方確實有一個升級腳本：
但是看上面日誌也就知道，它差不多等同於「import tensorflow.compat.v1 as tf」。真正要利用上 TF 2.0 的 Eager Exexution，還是得手動重寫。
API 接口，難以明了
Tensorflow 1.X 時代，靜態圖雖說上手稍微難了那麼一丟丟，但是這並不是什麼問題。既然入了機器學習的坑，這當然是能掌握的。
TensorFlow 1.X 不好用的主要原因，還在於 API 接口比較 混亂。
tf.nntf.layertf.keras.layerstf.contrib.layertf.contrib.slim說到底，它們都是基本的神經網絡層級，很多時候都是有重疊的。這種 API 上的冗餘，極大地降低了 TF 的生態質量。尤其是，很多官方教程，很多谷歌開源的模型代碼，都用的是 tf.contrib.slim 來寫模型。比如說 MobileNet 之類的經典模型，官方實現就是用 TF 第三方庫「contrib」中的一個模塊「slim」來寫的。
然後到了 TensorFlow 2.X，整個「contrib」庫都被放棄了。在 1.X 後期，各個教程使用的接口都不相同，我們又分不清楚哪個接口到底好，哪個到底差。由此引出來的，就不僅是很差的用戶體驗，同時還有性能上的差異。如果我們用 1.X 中的 tf.nn.rnn_cell 來做 LSTM，這也是沒問題的，只不過會特別慢。如果我們將運算子換成 LSTM，那麼無疑速度會提升很多。整個 TF 1.X，在 API 接口上，總是存在大量的坑，需要算法工程師特別注意。那麼 TensorFlow 2.X 呢？雖然說 TF 2.X 方向很明確，默認採用動態計算圖，大力推進 tf.keras 這樣的高級 API。這些都非常好，甚至用 Keras 寫模型比 PyTorch 還要精簡一些。但是別忘了 TF 傳統藝能是靜態計算圖，它天生就比 tf.keras 擁有更多的底層配置。這就會導致兩種割裂的代碼風格，一種是非常底層，使用 tf.function 等更一般的 API 構建模型，能進行各方面的定製化。另一種則非常抽象，使用 tf.keras 像搭積木一樣搭建模型，我們不用了解底層的架構如何搭建，只需要關注整體的設計流程即可。如果教程與 API 對兩種模式都分的清清楚楚還好，但問題在於，引入 Keras 卻讓 API 又變得更加混亂了。
TF 2.X 官方教程目前以 Keras API 為主。這其實和 1.X 的情況還是挺像的，同一個功能能由不同的 API 實現，但是不同 API 進行組合的時候，就會出問題。也就是說，如果我們混淆了 tf.keras 和底層 API，那麼這又是一個大坑。比如說使用 tf.keras，以 model = tf.keras.Sequential 的方式構建了模型。那麼訓練流程又該是什麼樣的？是直接用 model.fit() ，還是說用 with tf.GradientTape() as Tape 做更具體的定製？如果我們先自定義損失函數，那這樣用高級 API 定義的模型，又該怎麼修改？
TF 2.X 官方教程，有的以類繼承的新方式，以及更底層的 API 構建模型。本來還沒進入 TF 2.X 時代時，keras 與 tf 兩部分 API 互不影響，該用哪個就用哪個。但是現在，tf.keras 中的高級 API，與 tf 中的底層 API 經常需要混用，這樣的整合會讓開發者不知所措。與此同時，API 的割裂，也加大了開發者尋找教程的難度。因為除了「TF2.0」 這個關鍵字，同時還要弄清楚：這個文檔是關於 TF2.0 本身的，還是關於 tf.keras 的。教程文檔不管怎麼說，TensorFlow 都是第一大深度學習框架，GitHub 上高達 15.2 萬的 Star 量遠超其它框架。其實在 TF 早期使用靜態計算圖的時期，整體教程還是比較連貫的，靜態計算圖有一條完整的路線。而且我們還有 tensor2tensor 這樣的代碼庫，裡面的代碼質量還是非常不錯的。後來隨著深度學習成為主流，也就有了各種非官方教程，tf.contrib 模塊裡面的代碼也就越來越多。到了 TF 2.X，tf.keras 整合進去之後，相關的文檔還是比較少的，以至於整個指引文檔成了 Keras 和經典 TF 的混合。還是拿之前的例子來說，在官方文檔上，如果要做圖像識別，教程會告訴你用 tf.keras.Sequential() 組合不同的神經網絡層級，然後依次定義 model.compile() 與 model.fit()，然後深度學習模型就能訓練起來了。
同樣，如果要做圖像生成模型，那麼教程還是告訴你用 tf.keras.Sequential() 組合神經網絡層級，但接下來卻需要自己定義損失函數、最優化器、控制迭代梯度等等。@tf.function、tf.GradientTape() 等等新模塊，都會用上。
採用 @tf.function、tf.GradientTape() 等 TF 2.X 新特性的一個示例。除了這兩種，對於更複雜的模型，TF2.0 還有一套解決方案，即從 tf.keras.Model 繼承模型，重新實現 call 方法。總之官方文檔有多種解決方案，能處理相同的問題。這種高級 API 與底層 API 混合在一起的做法特別常見，因此很多時候會感覺 TF 2.X 的技術路線不是非常明確。此外，tf.keras 是個「大雜燴」，神經網絡層級、最優化器、損失函數、數據預處理 API 等等都包含在內。它要與 tf.nn、tf.train、tf.data 之類的 API 在相同層級，總感覺有點怪怪的。看上去底層 API 與高級 API 兩大類文檔，應該是平級的，這樣找起來比較好理解。
還有一點：速度按理來說，TensorFlow 的速度優化的應該還是可以的。而且本來各框架的速度差別就不是很明顯，所以速度上應該沒什麼問題。但是在做算法工程的時候，總會聽到周圍的朋友抱怨 TensorFlow 的速度還不如 PyTorch，抱怨 TF 2.0 儘管用上了動態計算圖，但速度還不如沒升級之前的 TF 1.X 代碼。這樣抱怨最大的可能性是，在做算法時選擇的 Kernel 不太對，或者計算流、數據流在某些地方存在瓶頸，甚至是某些訓練配置就根本錯了。所以說，速度方面，很可能是我們自己優化沒做到位。但是 TF 1.X 升級到 2.X 之後，速度真的會有差別嗎？筆者還真的做過非標準測試，如果使用升級腳本完成升級，同樣的代碼，兩者底層的計算子還真不一樣。速度上甚至 TF 1.X 略有優勢。在最初的 TF 1.X 代碼中，很多矩陣運算用的都是 tf.einsum()，靜態計算圖應該把它都轉化為了 MatMul 等計算子，整體推斷速度平均是 16ms。
相同代碼，在 TF 1.X 下的推斷速度。而如果根據 tf.compat.v1 升級代碼，相同的模型確實底層計算子都不太一樣。但沒想到的是，TF 2.X 採用了新的 Einsum 運算，速度好像並不佔優？
相同代碼，在 TF 2.X 下的推斷速度。小結最後，我們想說的是，作為AI工程師，選擇適合自己的深度學習框架需要認真考慮。雖說大家很多都用 TensorFlow，但維護起來真的有挺多坑要踩。而且一升級TF 2.X，就相當於換了個框架，十分麻煩，甚至不如根據業務情況重新考慮其他框架。放眼望去，國產的幾種框架也許會是不錯的選擇。它們至少都能經受得起業務的考驗，且與框架維護者交流起來也會比較方便。此外，一些國產新興框架，沒有歷史包袱，技術路線比較統一與確定。一般而言，只要訓練效率高、部署方便高效、代碼比較好維護，且能滿足業務所在領域的需求，那麼這就是個好框架。總而言之，適合自己的才是最好的。