經(jīng)驗：遷移、訓(xùn)練、部署一站式指南)
坦白說看到tensorflow這個標(biāo)題我第一反應(yīng)不是該從什么基礎(chǔ)概念講起而是想聊聊這兩年我自己從TensorFlow 1.x遷到2.x的真實經(jīng)歷。2019年那會兒我還在維護著一套用1.x寫的圖像分類系統(tǒng)每次訓(xùn)練啟動要寫一大堆tf.Session()模板代碼模型調(diào)試全靠print大法打印每個節(jié)點的輸出當(dāng)時就隱隱覺得這套東西遲早要變。果然等TensorFlow 2.x正式發(fā)布之后原來的寫法幾乎被全部推倒重來。對于已經(jīng)習(xí)慣1.x語法的人來說那段過渡期確實痛苦但對于新入門的朋友來說2.x的編程體驗其實比老版本舒服太多。這篇文章不打算面面俱到地羅列TensorFlow的每一個API那樣讀起來跟官方文檔沒什么區(qū)別。我只想從實際使用的角度出發(fā)把TensorFlow 2.x最核心的機制講清楚再分享一些我在數(shù)據(jù)管道、訓(xùn)練調(diào)參、模型部署各個環(huán)節(jié)里摸爬滾打攢下的經(jīng)驗。如果你正在糾結(jié)怎么入門TensorFlow或者已經(jīng)用了一段時間但總覺得沒過到點上又或者打算把訓(xùn)練好的模型真正推到生產(chǎn)環(huán)境那這篇內(nèi)容應(yīng)該能幫你省下不少彎路。1. 從TensorFlow 1.x到2.x的遷移我為什么放棄了大量舊代碼先說個我的真實經(jīng)歷。有個長期維護的商品識別項目底層用的是TensorFlow 1.14主要依靠tf.estimator框架管理分布式訓(xùn)練。當(dāng)時項目跑得很穩(wěn)定模型精度也已達(dá)到業(yè)務(wù)指標(biāo)所以在團隊討論要不要升級到2.x時我本能地抗拒。但后來一次新需求讓我徹底改變了想法新模型需要采用一種動態(tài)序列結(jié)構(gòu)在1.x里意味著要手動維護多個feature_column、寫復(fù)雜的model_fn同時還要處理大量tf.cond和循環(huán)控制——這套組合下來代碼量翻倍而且bug頻出。當(dāng)時我就意識到繼續(xù)守著1.x實際上是在用高昂的開發(fā)成本換一時的穩(wěn)定。1.1 1.x和2.x在思路上的天壤之別TensorFlow 1.x的核心邏輯是先建圖、后執(zhí)行。你要先通過一系列節(jié)點操作把計算圖搭建好然后在tf.Session里啟動會話執(zhí)行這個圖。這種方式在分布式部署上有優(yōu)勢但對普通開發(fā)者來說心智負(fù)擔(dān)太重。更麻煩的是你很難在中間某個節(jié)點直接打印結(jié)果必須把想要查看的張量也放進(jìn)sess.run()的fetches列表里調(diào)試一個單元至少要多寫三行代碼。2.x從根上把這種靜態(tài)圖優(yōu)先的策略改掉了默認(rèn)開啟即時執(zhí)行模式。也就是說你寫代碼時調(diào)用的每個TensorFlow算子會立刻返回真實數(shù)值而不是一個圖節(jié)點。這幾乎和寫普通Python代碼一樣直覺化對模型調(diào)試、快速原型驗證非常友好。很多人說TensorFlow 2.x變簡單了本質(zhì)就是因為這個語法層面的顛覆。另外一個關(guān)鍵變化是tf.keras被完全收編進(jìn)TensorFlow主庫。1.x時代我們往往要在keras和TensorFlow之間做選擇估值、優(yōu)化器、模型保存格式都各成一套。到了2.x官方直接建議所有人用tf.keras接口去搭建模型底層與TensorFlow緊密耦合既有高層API的簡潔又能在需要時隨時插入自定義的底層邏輯。1.2 遷移時真正要小心的地方如果你也是從1.x開始用、如今想遷移到2.x我的建議是別試圖逐行翻譯舊代碼那是事倍功半的路子。我建議把重心放在這幾件事上。第一拋棄tf.Session()相關(guān)的一切寫法。把原本包在with tf.Session()里的前向計算邏輯改寫成普通的Python函數(shù)返回你要的Tensor結(jié)果即可。第二把tf.placeholder改為函數(shù)的輸入?yún)?shù)。2.x里不再需要預(yù)先聲明占位的張量直接傳實參進(jìn)去動態(tài)形狀用None或者tf.TensorShape([None, ...])表示。第三tf.contrib目錄整體被移除里面很多擴展功能要么遷到tensorflow_addons要么被重構(gòu)成獨立模塊。項目里一旦出現(xiàn)tf.contrib相關(guān)引用編譯器會直接報錯你得逐個排查替換。我還遇到過一類隱蔽問題用舊代碼訓(xùn)練的checkpoint在新版本的TensorFlow里并不保證能原樣恢復(fù)。特別是用tf.estimator保存的模型目錄內(nèi)部結(jié)構(gòu)中帶有大量舊的命名空間想加載到2.x模型里最穩(wěn)妥的做法是只取出權(quán)重再用2.x代碼重建模型結(jié)構(gòu)后按變量名賦值。至于tf.layers與tf.keras.layers之間的對應(yīng)關(guān)系雖然很多名字長得差不多但底層的變量初始化方式和正則化策略會有差異遷移后建議做一次完整的小規(guī)模訓(xùn)練驗證而不是直接把舊權(quán)重丟進(jìn)去就完事。這里附一個我當(dāng)時對照用的API變更表幫助大家快速定位替換目標(biāo)TensorFlow 1.xTensorFlow 2.x 替代方案tf.Session()去掉Session直接執(zhí)行算子或使用tf.functiontf.placeholder函數(shù)參數(shù)tf.layers.densetf.keras.layers.Densetf.contrib.rnntf.keras.layers.LSTM或tensorflow_addonstf.get_variabletf.Variabletf.summary.scalartf.summary配合TensorBoard回調(diào)總的來說遷移這件事如果能提前規(guī)劃好路徑后面反而會輕松很多。我在遷完那套商品識別模型之后最大的感受是代碼量縮減了將近四成而且團隊成員新上手的速度快了一截因為新一代寫法跟平時寫Python的習(xí)慣幾乎一致。2. 張量、計算圖與自動微分TensorFlow運行機制的通俗拆解很多人覺得TensorFlow難學(xué)往往是卡在內(nèi)部機制上。其實拆開來看就是三件事張量怎么表示數(shù)據(jù)、計算圖怎么組織運算、自動微分怎么幫你求梯度。這三個東西搞懂了后面再遇到任何新API你都不會覺得跟魔法一樣。2.1 張量把數(shù)據(jù)當(dāng)作多維數(shù)組所謂張量你可以直接把它理解為帶著數(shù)據(jù)類型和多維形狀的數(shù)組。零階張量就是標(biāo)量比如一個普通數(shù)字一階張量類似一維列表二階張量就是人們常說的矩陣再往上就是更高維度的數(shù)據(jù)。TensorFlow里幾乎所有的數(shù)據(jù)傳遞都在張量層面完成包括數(shù)值計算、模型權(quán)重的更新、神經(jīng)網(wǎng)絡(luò)中間層輸出的保存等等底層全是張量操作。用代碼創(chuàng)建張量非常直白TensorFlow提供了一整套生成工具import tensorflow as tf # 從一個Python列表創(chuàng)建一階張量 a tf.constant([1, 2, 3], dtypetf.int32) # 創(chuàng)建全是0的二階張量 b tf.zeros(shape(2, 3), dtypetf.float32) # 創(chuàng)建在0到1之間均勻分布的隨機張量 c tf.random.uniform(shape(4, 4), minval0.0, maxval1.0) print(a.shape, b.shape, c.shape)這里我很想說的一點是剛?cè)腴T時務(wù)必要搞清楚shape和dtype這兩個屬性。TensorFlow對張量形狀的匹配檢查非常嚴(yán)格一個形狀為(32, 10)的矩陣想和另一個形狀為(32, 20)的矩陣做矩陣乘法會直接拋出維度不匹配的異常。而dtype更像是一個電路的規(guī)范如果你讓一個float32張量與int32張量直接相加TensorFlow會在某些情況下主動幫你報錯而不是做隱式類型轉(zhuǎn)換。這一點跟Python原生的什么都幫你自動轉(zhuǎn)不太一樣習(xí)慣之后其實反而能減少很多隱晦bug。2.2 計算圖把計算當(dāng)作生產(chǎn)線雖然即時執(zhí)行模式是2.x的默認(rèn)交互方式但TensorFlow在底層依然保留了一套基于計算圖的加速機制。尤其在訓(xùn)練大批量數(shù)據(jù)或部署到生產(chǎn)環(huán)境時靜態(tài)圖帶來的性能收益相當(dāng)明顯。這里的思路用生產(chǎn)線類比你把每個數(shù)學(xué)步驟看成一臺專用機器輸入數(shù)據(jù)從一臺機器流到下一臺機器最后從出口出成品。因為整個流程是預(yù)先布置好的執(zhí)行時不需要每次重新解釋語句底層的編譯器還可以對調(diào)度、內(nèi)存分配做優(yōu)化。在2.x中你不需要手動構(gòu)建整張圖只需要在普通Python函數(shù)上加上tf.function裝飾器TensorFlow就會把這個函數(shù)里的TensorFlow操作圖化處理。舉個例子tf.function def compute_linear(w, x, b): y tf.matmul(w, x) b return tf.nn.relu(y) w tf.random.normal((16, 8)) x tf.random.normal((8, 1)) b tf.zeros((16, 1)) result compute_linear(w, x, b)這看起來就像一個普通函數(shù)調(diào)用但TensorFlow在首次調(diào)用時會把內(nèi)部操作編譯為圖后續(xù)再次調(diào)用就能復(fù)用這份編譯結(jié)果。這里有個關(guān)鍵點要特別提醒tf.function裝飾的函數(shù)內(nèi)部如果使用了Python的if、while等控制流TensorFlow會利用一種叫自動控制流轉(zhuǎn)換的機制把它們改成圖的控制流節(jié)點但有時候這種轉(zhuǎn)換會跟外部的Python狀態(tài)產(chǎn)生微妙的不一致。比如在函數(shù)里直接打印一個Tensor對象你會看到它是抽象符號而非具體值除非調(diào)用print(張量.numpy())之類的顯式操作。我也見過有人過于依賴tf.function結(jié)果把數(shù)據(jù)處理邏輯也寫進(jìn)去反而讓調(diào)試變得很麻煩。我的實踐體會是能用即時執(zhí)行模式寫清楚邏輯的時候不用急著做圖化訓(xùn)練速度不夠再針對性優(yōu)化。2.3 自動微分用GradientTape記錄梯度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的核心是反向傳播也就是求損失函數(shù)對每個參數(shù)的偏導(dǎo)數(shù)。在TensorFlow 2.x中你用tf.GradientTape這個環(huán)境管理器就能非常方便地獲得梯度。思路是把需要求導(dǎo)的計算過程包進(jìn)with tf.GradientTape() as tape:里運行時TensorFlow會記錄所有涉及tf.Variable的張量操作之后調(diào)用tape.gradient()就能根據(jù)目標(biāo)張量和變量列表算出梯度。看一個簡單的線性回歸示例import tensorflow as tf # 定義參數(shù) w tf.Variable(tf.random.normal((1,)), namew) b tf.Variable(tf.zeros((1,)), nameb) # 輸入數(shù)據(jù)這里取示例值 x tf.constant([1.0, 2.0, 3.0, 4.0]) y tf.constant([2.0, 4.0, 6.0, 8.0]) # 訓(xùn)練若干輪 for epoch in range(50): with tf.GradientTape() as tape: y_pred w * x b loss tf.reduce_mean(tf.square(y_pred - y)) grads tape.gradient(loss, [w, b]) # 手動完成一個梯度下降步驟 w.assign_sub(0.01 * grads[0]) b.assign_sub(0.01 * grads[1]) print(fw: {w.numpy()}, b: {b.numpy()})這里tape.gradient(loss, [w, b])返回的梯度列表順序和傳入變量列表完全一致。實際訓(xùn)練中優(yōu)化器比如tf.keras.optimizers.Adam的apply_gradients方法就可以直接接收這個梯度列表完成參數(shù)更新。這個機制看起來很基礎(chǔ)但它正是整條深度學(xué)習(xí)訓(xùn)練鏈路的基石。理解它之后你再去看Keras的Model.fit就不會覺得內(nèi)部有多神秘了——它其實也走的是這一套求梯度和更新參數(shù)的過程無非外面包了更多調(diào)度和狀態(tài)管理。3. 讓數(shù)據(jù)管道跑起來tf.data的高效使用心得幾乎每個用過TensorFlow的人最后都會得出同一個結(jié)論模型架構(gòu)只是深度學(xué)習(xí)系統(tǒng)的一部分如果數(shù)據(jù)管道設(shè)計得不好再好的模型也會因為喂數(shù)據(jù)太慢而在訓(xùn)練時白白等待。我自己就在這個環(huán)節(jié)踩過不少坑尤其當(dāng)數(shù)據(jù)量上了幾十萬張圖片之后每一輪訓(xùn)練里GPU有一半時間在空轉(zhuǎn)的情況非常常見。3.1 用tf.data構(gòu)建數(shù)據(jù)集的基本姿勢最基礎(chǔ)的數(shù)據(jù)集構(gòu)造方式是從內(nèi)存數(shù)組或文件路徑批量創(chuàng)建。tf.data.Dataset.from_tensor_slices常用于把圖片路徑和標(biāo)簽組合成訓(xùn)練集import tensorflow as tf data_paths [image_1.jpg, image_2.jpg, image_3.jpg] labels [0, 1, 0] dataset tf.data.Dataset.from_tensor_slices((data_paths, labels)) # 之后你可以按需對dataset做變換 dataset dataset.map(load_and_preprocess_image) dataset dataset.shuffle(1000) dataset dataset.batch(32)如果數(shù)據(jù)和標(biāo)簽都存在于內(nèi)存這種方式很直接。不過對于超大訓(xùn)練集建議先把數(shù)據(jù)寫入TFRecord格式。TFRecord是TensorFlow自定義的一種二進(jìn)制存儲格式對讀取速度友好而且可以借助tf.io.parse_single_example在數(shù)據(jù)流中進(jìn)行解碼。缺點是序列化和解析代碼寫起來有點繁瑣所以大多數(shù)小項目并不需要一上來就上TFRecord等內(nèi)存放得下的時候簡單的文件路徑列表加map變換就夠了。3.2 數(shù)據(jù)Pipeline的三大性能加速器我剛開始訓(xùn)練圖像模型時一個epoch要跑接近二十分鐘后來優(yōu)化了數(shù)據(jù)管道直接壓到了七分鐘。收益主要來自三件事并行map、預(yù)取和緩存。并行map很好理解dataset.map(load_image, num_parallel_callstf.data.AUTOTUNE)讓數(shù)據(jù)預(yù)處理解碼圖片、縮放、歸一化在多個線程里并行進(jìn)行。tf.data.AUTOTUNE是TensorFlow根據(jù)實際環(huán)境自動調(diào)整的并行度數(shù)值省去手動調(diào)參的麻煩。預(yù)取則是把下一個batch的準(zhǔn)備工作和當(dāng)前batch的計算重疊起來。調(diào)用dataset.prefetch(buffer_sizetf.data.AUTOTUNE)之后系統(tǒng)會在GPU還在訓(xùn)練當(dāng)前batch時提前準(zhǔn)備后續(xù)數(shù)據(jù)這能直接把訓(xùn)練間隙的等待時間壓到最低。我用下來的感受是prefetch是性價比最高的一個操作哪怕只是加一行延遲也能明顯下降。緩存的作用就更直接了dataset.cache()可以把第一個epoch計算過的預(yù)處理結(jié)果緩存到內(nèi)存或磁盤文件中。如果數(shù)據(jù)量不太大緩存到內(nèi)存后后續(xù)每個epoch幾乎不花預(yù)處理時間但要注意如果shuffle操作寫在cache之前緩存的就是打亂后的順序這跟你預(yù)期的隨機性可能會有出入建議把cache()放在shuffle()前面。這里我整理了一個典型的數(shù)據(jù)管道順序供你參考dataset dataset.cache() dataset dataset.shuffle(buffer_size5000) dataset dataset.map(preprocess, num_parallel_callstf.data.AUTOTUNE) dataset dataset.batch(32) dataset dataset.prefetch(tf.data.AUTOTUNE)3.3 順序問題為什么shuffle和repeat要放在恰當(dāng)位置很多人沒留意shuffle、repeat與batch之間的順序關(guān)系結(jié)果訓(xùn)練出的模型效果總差一口氣。我對原則的理解是如果希望每個epoch里都看到不同批次的組合shuffle應(yīng)該在batch之前執(zhí)行如果希望數(shù)據(jù)一直循環(huán)供應(yīng)repeat可以放在batch之前或之后取決于你想要的epoch語義。一般習(xí)慣是dataset dataset.shuffle(...).batch(...).repeat(...)這樣每個epoch內(nèi)數(shù)據(jù)是打亂的而且epoch之間數(shù)據(jù)分布也會不同。還有一個容易被忽略的問題在Dataset中如果先執(zhí)行batch再執(zhí)行shuffleshuffle的粒度就變成了隨機重排batch而不是隨機重排樣本。對于大多數(shù)監(jiān)督學(xué)習(xí)任務(wù)來說這樣會讓同一個batch里的樣本彼此高度相關(guān)訓(xùn)練初期模型的loss曲線會明顯抖動甚至影響收斂。所以我的建議是sample級別的shuffle要放在batch之前。4. 訓(xùn)練過程的可控與調(diào)優(yōu)從損失不降說起模型訓(xùn)練調(diào)試是TensorFlow使用中耗時最多的環(huán)節(jié)之一。即便代碼語法完全正確模型也常常不按預(yù)期收斂。我自己經(jīng)歷過一次印象深刻的調(diào)優(yōu)過程一個文本分類模型前五個epoch的loss一直卡在2.0左右不下去剛開始我懷疑是學(xué)習(xí)率太高導(dǎo)致震蕩于是降到1e-4結(jié)果loss紋絲不動接著懷疑模型結(jié)構(gòu)有錯又反復(fù)檢查了好幾遍仍沒發(fā)現(xiàn)問題。最后發(fā)現(xiàn)問題其實出在標(biāo)簽預(yù)處理上——我把多標(biāo)簽任務(wù)的標(biāo)簽喂成了一個形狀與模型輸出不匹配的張量損失函數(shù)的掩碼邏輯把這個不匹配悄悄掩蓋了反而導(dǎo)致梯度一直在原地打轉(zhuǎn)。4.1 損失不降的排查順序如果你也遇到loss異常平穩(wěn)或抖動不降的情況我建議按下面這個順序排查能省不少時間。先檢查數(shù)據(jù)與標(biāo)簽。打印一個batch的數(shù)據(jù)分布看看類別是否均衡標(biāo)簽是否有大量錯位歸一化是否做到位。尤其是文本序列任務(wù)padding掉的token如果直接參與損失計算也會拖住收斂這時可以借助tf.keras.layers.Masking或自定義損失函數(shù)中的mask處理。再檢查網(wǎng)絡(luò)的輸入端。輸入數(shù)據(jù)形狀與模型首層期待的shape是否一致比如卷積網(wǎng)絡(luò)輸入是不是(batch, height, width, channels)序列模型輸入是不是(batch, seq_len, features)。很多時候shape完全匹配但值域不對比如圖像沒有歸一化到[0,1]也會讓訓(xùn)練初期不穩(wěn)定。然后是優(yōu)化器與學(xué)習(xí)率。建議用tf.keras.optimizers.Adam作為默認(rèn)優(yōu)化器起步學(xué)習(xí)率先從1e-3或者3e-4這種量級試起。如果loss出現(xiàn)明顯震蕩減小學(xué)習(xí)率如果loss一直不動而不是強烈震蕩那問題多半不在學(xué)習(xí)率而在數(shù)據(jù)或梯度本身。最后可以用梯度檢查確認(rèn)模型梯度是否在正常范圍。tape.gradient得到的梯度如果出現(xiàn)大量NaN或者接近0的數(shù)值通常意味著網(wǎng)絡(luò)前面有數(shù)值不穩(wěn)定或者梯度消失。這時候圍繞網(wǎng)絡(luò)初始化、激活函數(shù)類型去排查比盲目調(diào)學(xué)習(xí)率更有用。4.2 學(xué)習(xí)率策略不是越傳越好在訓(xùn)練早期我習(xí)慣讓模型在一個較小學(xué)習(xí)率下熱身再慢慢上升到目標(biāo)學(xué)習(xí)率最后通過衰減逐步降低。TensorBoard里可以清楚看到這些策略對訓(xùn)練曲線的影響。tf.keras.optimizers.schedules提供了多種現(xiàn)成的學(xué)習(xí)率調(diào)度方式比如ExponentialDecay比較經(jīng)典CosineDecay在中大型任務(wù)上也常見。一個典型的用法是import tensorflow as tf lr_schedule tf.keras.optimizers.schedules.ExponentialDecay( initial_learning_rate1e-3, decay_steps1000, decay_rate0.96, staircaseTrue ) optimizer tf.keras.optimizers.Adam(learning_ratelr_schedule)在實際項目中我發(fā)現(xiàn)很多情況下恒定學(xué)習(xí)率加一個足夠長的預(yù)熱期效果就已經(jīng)可以了復(fù)雜的調(diào)度方案并不一定會帶來顯著提升。學(xué)習(xí)率這個超參數(shù)高度依賴數(shù)據(jù)規(guī)模和任務(wù)復(fù)雜度與其照搬別人的參數(shù)不如小規(guī)模實驗幾次然后觀察loss在訓(xùn)練過程中如何變化。4.3 過擬合問題不能靠早停一招解決很多新手一發(fā)現(xiàn)驗證集loss開始回升就立刻加上早停機制。這個做法短期有效但掩蓋了模型泛化能力不足的本質(zhì)。我更常用的一套組合拳是數(shù)據(jù)增強、權(quán)重正則化和Dropout。數(shù)據(jù)增強在圖像任務(wù)上非常好用。TensorFlow有大量內(nèi)置變換比如隨機裁剪、旋轉(zhuǎn)、翻轉(zhuǎn)、亮度對比度調(diào)整等直接在數(shù)據(jù)管道里實現(xiàn)即可。文本任務(wù)則可以用隨機的同義詞替換或弱化高頻詞等策略。權(quán)重正則化通過kernel_regularizer參數(shù)就能實現(xiàn)常見的有L1和L2正則。Dropout一般加在全連接層之后它本質(zhì)是在訓(xùn)練時隨機失活一部分神經(jīng)元迫使網(wǎng)絡(luò)學(xué)到冗余表示從而提高泛化能力。把這些手段組合起來模型即使訓(xùn)練更久也不會過早地在驗證集上惡化。5. 從訓(xùn)練到生產(chǎn)模型導(dǎo)出、量化與部署落地訓(xùn)練管道跑通只是第一步真正讓TensorFlow模型產(chǎn)生業(yè)務(wù)價值的是把模型部署到線上環(huán)境。大家在本地Jupyter里model.predict跑得順溜但一推到生產(chǎn)就各種不兼容這種情況我見得太多了。5.1 用SavedModel統(tǒng)一導(dǎo)出流程TensorFlow 2.x推薦的模型導(dǎo)出格式是SavedModel。調(diào)用model.save(my_model_dir)即可把模型的網(wǎng)絡(luò)結(jié)構(gòu)、權(quán)重與推理函數(shù)打包到一個目錄里。之后用tf.keras.models.load_model(my_model_dir)就能把模型完整恢復(fù)出來。SavedModel的好處是結(jié)構(gòu)自包含不管是用Keras定義的模型還是底層自定義的tf.Module都能統(tǒng)一導(dǎo)出而且可以同時保存多種推理簽名方便不同調(diào)用方使用。如果你的模型除了常規(guī)預(yù)測之外還需要額外的預(yù)處理或后處理邏輯可以把這些邏輯也封裝成一個函數(shù)。這里給個比較通用的服務(wù)端部署思路先用原生TensorFlow定義推理函數(shù)然后用tf.function把它和模型打包成SavedModel再借助模型服務(wù)框架對外提供HTTP或gRPC接口。import tensorflow as tf # 假設(shè)model是一個訓(xùn)練好的keras模型 model tf.keras.models.load_model(my_model_dir) tf.function(input_signature[tf.TensorSpec(shape[None], dtypetf.string)]) def serve_fn(serialized_example): parsed tf.io.parse_tensor(serialized_example, out_typetf.float32) parsed tf.reshape(parsed, [-1, 224, 224, 3]) logits model(parsed) return {probability: tf.nn.softmax(logits)} tf.saved_model.save(model, exported_model, signatures{serving_default: serve_fn})這樣部署出來的服務(wù)調(diào)用方就不需要關(guān)心模型內(nèi)部細(xì)節(jié)只需要把輸入數(shù)據(jù)按約定格式傳進(jìn)來即可。簽名機制就像給模型開了一扇指定的門不同的服務(wù)方可以訪問不同的入口。5.2 輕量化TensorFlow Lite與量化如果模型最終要跑在移動端或嵌入式設(shè)備上直接用SavedModel肯定不合適。TensorFlow Lite是面向移動端和嵌入式場景的輕量級推理框架它把模型轉(zhuǎn)換為FlatBuffer格式針對移動設(shè)備做了算子精簡和內(nèi)存優(yōu)化。轉(zhuǎn)換方式非常簡單converter tf.lite.TFLiteConverter.from_saved_model(exported_model) tflite_model converter.convert() with open(model.tflite, wb) as f: f.write(tflite_model)如果你對模型體積還比較在意可以開啟量化。量化簡單理解就是降低權(quán)重和激活值的數(shù)值精度比如從32位浮點降為8位整數(shù)換來體積減少和運行速度提升代價是精度的輕微下降。TensorFlow Lite的量化方式有動態(tài)范圍量化、完整整數(shù)量化和浮點16量化等要在部署前對比一下量化后模型與原始模型在驗證集上的指標(biāo)差多少確認(rèn)下降幅度可接受再上線。5.3 服務(wù)化部署時常見的隱藏問題有兩點我在部署階段反復(fù)踩坑特別提醒。第一訓(xùn)練階段使用tf.data處理原始數(shù)據(jù)和部署階段接收請求數(shù)據(jù)之間的預(yù)處理邏輯要保持一致。很多模型線上效果不如預(yù)期往往就是因為本地測試時走了數(shù)據(jù)管道的完整預(yù)處理而線上服務(wù)直接拿原始字符串或二進(jìn)制數(shù)據(jù)喂給模型導(dǎo)致輸入分布產(chǎn)生了偏移。解決辦法就是把預(yù)處理邏輯完整封裝到serve_fn里不管原始數(shù)據(jù)是什么形式統(tǒng)一在模型入口完成轉(zhuǎn)換。第二不同TensorFlow版本導(dǎo)出的SavedModel在更高或更低版本運行時可能產(chǎn)生算子兼容性問題。生產(chǎn)環(huán)境建議鎖定一個長期維護的TensorFlow版本更新版本前先用代表性樣本做一輪推理一致性校驗確保輸出差異在可接受范圍內(nèi)。這一點對涉及金融風(fēng)控、醫(yī)療輔助這類對數(shù)值穩(wěn)定性要求高的場景尤其重要。6. 我踩過的TensorFlow深坑與繞行方案TensorFlow作為一個生態(tài)龐大的框架雖然好用但并不意味著沒有坑。有些問題官網(wǎng)文檔里寫得并不顯眼只有自己撞上去才知道有多疼。我把印象最深的一些經(jīng)驗整理在這里算是給后來的朋友一些繞行提示。6.1 版本與Python環(huán)境的匹配問題TensorFlow對Python版本有明確的兼容范圍比如某些較老版本在Python 3.10上可能直接編譯失敗。最省心的做法是用虛擬環(huán)境單獨安裝TensorFlow并且按官方文檔給定的Python版本范圍來創(chuàng)建環(huán)境。python -m venv tf_env source tf_env/bin/activate pip install tensorflow裝好之后可以用python -c import tensorflow as tf; print(tf.__version__)確認(rèn)安裝版本。我就遇到過因為錯誤使用舊版本安裝指令導(dǎo)致裝上了CPU版TensorFlow結(jié)果GPU始終無法識別后來排查才發(fā)現(xiàn)是包名寫錯了?，F(xiàn)在GPU版本不需要單獨安裝官方主包已經(jīng)同時支持CPU和GPU只要機器環(huán)境具備CUDA和cuDNN即可。6.2 GPU顯存不釋放與分配策略訓(xùn)練一個模型之后再做推理時發(fā)現(xiàn)GPU顯存仍然被占用這是新手特別容易困惑的問題。TensorFlow默認(rèn)會預(yù)先申請盡可能多的顯存而不是按需分配。如果不設(shè)置顯存增長策略即使你的模型只需要2GB顯存TensorFlow也可能先把整個顯存吃滿。我這里說的按需可以用一行代碼配置gpus tf.config.experimental.list_physical_devices(GPU) if gpus: try: tf.config.experimental.set_memory_growth(gpus[0], True) except RuntimeError as e: print(e)set_memory_growth設(shè)置為True后顯存會隨著實際需求動態(tài)分配不會一上來就占滿。如果確實希望固定使用某塊卡的固定顯存量用tf.config.set_logical_device_configuration配合tf.config.LogicalDeviceConfiguration(memory_limit...)更靈活。6.3 多線程和進(jìn)程的程序問題TensorFlow內(nèi)部有很多線程池來控制并行算子執(zhí)行。如果你在同一個Python進(jìn)程里同時加載多個模型并頻繁新建會話容易遇到默認(rèn)線程池配置被覆蓋的問題表現(xiàn)為推理耗時突然變長。我后來的習(xí)慣是盡量復(fù)用同一個模型實例來做多次推理而不是每次預(yù)測前都重新加載一次真有多個模型需要并發(fā)處理時用多進(jìn)程而非多線程來隔離資源。另一個和線程相關(guān)的典型問題是numpy與TensorFlow混合使用時不注意張量同步點。比如在繼承tf.keras.Model的自定義call方法里如果你調(diào)用了.numpy()方法把Tensor轉(zhuǎn)成numpy數(shù)組雖然能得到數(shù)值但這個操作會打斷正在構(gòu)建的自動微分鏈路導(dǎo)致tape.gradient無法正常工作。所以在訓(xùn)練路徑里盡量別在模型內(nèi)層使用.numpy()只在外部檢查結(jié)果或做數(shù)據(jù)IO時使用。6.4 隨機種子問題帶來的不可復(fù)現(xiàn)我早年復(fù)現(xiàn)實驗結(jié)果時曾經(jīng)因為忘記設(shè)置隨機種子而白白浪費了一周時間去找模型性能波動的原因。TensorFlow的隨機性涉及多個層級Python內(nèi)置的random模塊、numpy的RandomState、TensorFlow自身的圖隨機算子。想要最大程度地固定隨機行為需要三層同時設(shè)import os import random import numpy as np import tensorflow as tf os.environ[PYTHONHASHSEED] 0 random.seed(42) np.random.seed(42) tf.random.set_seed(42)不過要說明的是即使都設(shè)了種子使用GPU并行算子時某些情況下仍然可能存在細(xì)微的非確定性輸出。如果科研或工程上對可復(fù)現(xiàn)性要求很高可以考慮限制部分算子使用CPU執(zhí)行或者在代碼里記錄關(guān)鍵變量的hash值方便對照。寫在最后的幾句真心話如果讓我給剛接觸TensorFlow的人提一個最實用的建議那一定是先沉住氣把張量形狀、dtype和各種常見的張量操作練熟再去看模型結(jié)構(gòu)。TensorFlow的語法體系非常龐雜但絕大多數(shù)操作歸根結(jié)底都是圍繞張量展開的。你會發(fā)現(xiàn)一旦具備了這種形狀思維讀官方文檔就像查字典看到某個API名字就能大致猜到它的輸入輸出形態(tài)。我自己在實際操作中還養(yǎng)成了一個習(xí)慣任何新方案上線前先用小數(shù)據(jù)集跑通一個最小實驗再逐步放大。這個習(xí)慣幫我在不少項目中及時發(fā)現(xiàn)了版本兼容和數(shù)據(jù)處理上的隱藏問題成本遠(yuǎn)比直接跑全量訓(xùn)練低得多。TensorFlow的生態(tài)仍在持續(xù)演進(jìn)新的API和最佳實踐層出不窮但核心思想是穩(wěn)定的。把這個框架用好靠的從來不是刷多少教程而是在一次次訓(xùn)練、調(diào)優(yōu)和部署的過程中積累起來的手感。 TensorFlow 2.x實戰(zhàn)經(jīng)驗：遷移、訓(xùn)練、部署一站式指南
坦白說看到tensorflow這個標(biāo)題我第一反應(yīng)不是該從什么基礎(chǔ)概念講起而是想聊聊這兩年我自己從TensorFlow 1.x遷到2.x的真實經(jīng)歷。2019年那會兒我還在維護著一套用1.x寫的圖像分類系統(tǒng)每次訓(xùn)練啟動要寫一大堆tf.Session()模板代碼模型調(diào)試全靠print大法打印每個節(jié)點的輸出當(dāng)時就隱隱覺得這套東西遲早要變。果然等TensorFlow 2.x正式發(fā)布之后原來的寫法幾乎被全部推倒重來。對于已經(jīng)習(xí)慣1.x語法的人來說那段過渡期確實痛苦但對于新入門的朋友來說2.x的編程體驗其實比老版本舒服太多。這篇文章不打算面面俱到地羅列TensorFlow的每一個API那樣讀起來跟官方文檔沒什么區(qū)別。我只想從實際使用的角度出發(fā)把TensorFlow 2.x最核心的機制講清楚再分享一些我在數(shù)據(jù)管道、訓(xùn)練調(diào)參、模型部署各個環(huán)節(jié)里摸爬滾打攢下的經(jīng)驗。如果你正在糾結(jié)怎么入門TensorFlow或者已經(jīng)用了一段時間但總覺得沒過到點上又或者打算把訓(xùn)練好的模型真正推到生產(chǎn)環(huán)境那這篇內(nèi)容應(yīng)該能幫你省下不少彎路。1. 從TensorFlow 1.x到2.x的遷移我為什么放棄了大量舊代碼先說個我的真實經(jīng)歷。有個長期維護的商品識別項目底層用的是TensorFlow 1.14主要依靠tf.estimator框架管理分布式訓(xùn)練。當(dāng)時項目跑得很穩(wěn)定模型精度也已達(dá)到業(yè)務(wù)指標(biāo)所以在團隊討論要不要升級到2.x時我本能地抗拒。但后來一次新需求讓我徹底改變了想法新模型需要采用一種動態(tài)序列結(jié)構(gòu)在1.x里意味著要手動維護多個feature_column、寫復(fù)雜的model_fn同時還要處理大量tf.cond和循環(huán)控制——這套組合下來代碼量翻倍而且bug頻出。當(dāng)時我就意識到繼續(xù)守著1.x實際上是在用高昂的開發(fā)成本換一時的穩(wěn)定。1.1 1.x和2.x在思路上的天壤之別TensorFlow 1.x的核心邏輯是先建圖、后執(zhí)行。你要先通過一系列節(jié)點操作把計算圖搭建好然后在tf.Session里啟動會話執(zhí)行這個圖。這種方式在分布式部署上有優(yōu)勢但對普通開發(fā)者來說心智負(fù)擔(dān)太重。更麻煩的是你很難在中間某個節(jié)點直接打印結(jié)果必須把想要查看的張量也放進(jìn)sess.run()的fetches列表里調(diào)試一個單元至少要多寫三行代碼。2.x從根上把這種靜態(tài)圖優(yōu)先的策略改掉了默認(rèn)開啟即時執(zhí)行模式。也就是說你寫代碼時調(diào)用的每個TensorFlow算子會立刻返回真實數(shù)值而不是一個圖節(jié)點。這幾乎和寫普通Python代碼一樣直覺化對模型調(diào)試、快速原型驗證非常友好。很多人說TensorFlow 2.x變簡單了本質(zhì)就是因為這個語法層面的顛覆。另外一個關(guān)鍵變化是tf.keras被完全收編進(jìn)TensorFlow主庫。1.x時代我們往往要在keras和TensorFlow之間做選擇估值、優(yōu)化器、模型保存格式都各成一套。到了2.x官方直接建議所有人用tf.keras接口去搭建模型底層與TensorFlow緊密耦合既有高層API的簡潔又能在需要時隨時插入自定義的底層邏輯。1.2 遷移時真正要小心的地方如果你也是從1.x開始用、如今想遷移到2.x我的建議是別試圖逐行翻譯舊代碼那是事倍功半的路子。我建議把重心放在這幾件事上。第一拋棄tf.Session()相關(guān)的一切寫法。把原本包在with tf.Session()里的前向計算邏輯改寫成普通的Python函數(shù)返回你要的Tensor結(jié)果即可。第二把tf.placeholder改為函數(shù)的輸入?yún)?shù)。2.x里不再需要預(yù)先聲明占位的張量直接傳實參進(jìn)去動態(tài)形狀用None或者tf.TensorShape([None, ...])表示。第三tf.contrib目錄整體被移除里面很多擴展功能要么遷到tensorflow_addons要么被重構(gòu)成獨立模塊。項目里一旦出現(xiàn)tf.contrib相關(guān)引用編譯器會直接報錯你得逐個排查替換。我還遇到過一類隱蔽問題用舊代碼訓(xùn)練的checkpoint在新版本的TensorFlow里并不保證能原樣恢復(fù)。特別是用tf.estimator保存的模型目錄內(nèi)部結(jié)構(gòu)中帶有大量舊的命名空間想加載到2.x模型里最穩(wěn)妥的做法是只取出權(quán)重再用2.x代碼重建模型結(jié)構(gòu)后按變量名賦值。至于tf.layers與tf.keras.layers之間的對應(yīng)關(guān)系雖然很多名字長得差不多但底層的變量初始化方式和正則化策略會有差異遷移后建議做一次完整的小規(guī)模訓(xùn)練驗證而不是直接把舊權(quán)重丟進(jìn)去就完事。這里附一個我當(dāng)時對照用的API變更表幫助大家快速定位替換目標(biāo)TensorFlow 1.xTensorFlow 2.x 替代方案tf.Session()去掉Session直接執(zhí)行算子或使用tf.functiontf.placeholder函數(shù)參數(shù)tf.layers.densetf.keras.layers.Densetf.contrib.rnntf.keras.layers.LSTM或tensorflow_addonstf.get_variabletf.Variabletf.summary.scalartf.summary配合TensorBoard回調(diào)總的來說遷移這件事如果能提前規(guī)劃好路徑后面反而會輕松很多。我在遷完那套商品識別模型之后最大的感受是代碼量縮減了將近四成而且團隊成員新上手的速度快了一截因為新一代寫法跟平時寫Python的習(xí)慣幾乎一致。2. 張量、計算圖與自動微分TensorFlow運行機制的通俗拆解很多人覺得TensorFlow難學(xué)往往是卡在內(nèi)部機制上。其實拆開來看就是三件事張量怎么表示數(shù)據(jù)、計算圖怎么組織運算、自動微分怎么幫你求梯度。這三個東西搞懂了后面再遇到任何新API你都不會覺得跟魔法一樣。2.1 張量把數(shù)據(jù)當(dāng)作多維數(shù)組所謂張量你可以直接把它理解為帶著數(shù)據(jù)類型和多維形狀的數(shù)組。零階張量就是標(biāo)量比如一個普通數(shù)字一階張量類似一維列表二階張量就是人們常說的矩陣再往上就是更高維度的數(shù)據(jù)。TensorFlow里幾乎所有的數(shù)據(jù)傳遞都在張量層面完成包括數(shù)值計算、模型權(quán)重的更新、神經(jīng)網(wǎng)絡(luò)中間層輸出的保存等等底層全是張量操作。用代碼創(chuàng)建張量非常直白TensorFlow提供了一整套生成工具import tensorflow as tf # 從一個Python列表創(chuàng)建一階張量 a tf.constant([1, 2, 3], dtypetf.int32) # 創(chuàng)建全是0的二階張量 b tf.zeros(shape(2, 3), dtypetf.float32) # 創(chuàng)建在0到1之間均勻分布的隨機張量 c tf.random.uniform(shape(4, 4), minval0.0, maxval1.0) print(a.shape, b.shape, c.shape)這里我很想說的一點是剛?cè)腴T時務(wù)必要搞清楚shape和dtype這兩個屬性。TensorFlow對張量形狀的匹配檢查非常嚴(yán)格一個形狀為(32, 10)的矩陣想和另一個形狀為(32, 20)的矩陣做矩陣乘法會直接拋出維度不匹配的異常。而dtype更像是一個電路的規(guī)范如果你讓一個float32張量與int32張量直接相加TensorFlow會在某些情況下主動幫你報錯而不是做隱式類型轉(zhuǎn)換。這一點跟Python原生的什么都幫你自動轉(zhuǎn)不太一樣習(xí)慣之后其實反而能減少很多隱晦bug。2.2 計算圖把計算當(dāng)作生產(chǎn)線雖然即時執(zhí)行模式是2.x的默認(rèn)交互方式但TensorFlow在底層依然保留了一套基于計算圖的加速機制。尤其在訓(xùn)練大批量數(shù)據(jù)或部署到生產(chǎn)環(huán)境時靜態(tài)圖帶來的性能收益相當(dāng)明顯。這里的思路用生產(chǎn)線類比你把每個數(shù)學(xué)步驟看成一臺專用機器輸入數(shù)據(jù)從一臺機器流到下一臺機器最后從出口出成品。因為整個流程是預(yù)先布置好的執(zhí)行時不需要每次重新解釋語句底層的編譯器還可以對調(diào)度、內(nèi)存分配做優(yōu)化。在2.x中你不需要手動構(gòu)建整張圖只需要在普通Python函數(shù)上加上tf.function裝飾器TensorFlow就會把這個函數(shù)里的TensorFlow操作圖化處理。舉個例子tf.function def compute_linear(w, x, b): y tf.matmul(w, x) b return tf.nn.relu(y) w tf.random.normal((16, 8)) x tf.random.normal((8, 1)) b tf.zeros((16, 1)) result compute_linear(w, x, b)這看起來就像一個普通函數(shù)調(diào)用但TensorFlow在首次調(diào)用時會把內(nèi)部操作編譯為圖后續(xù)再次調(diào)用就能復(fù)用這份編譯結(jié)果。這里有個關(guān)鍵點要特別提醒tf.function裝飾的函數(shù)內(nèi)部如果使用了Python的if、while等控制流TensorFlow會利用一種叫自動控制流轉(zhuǎn)換的機制把它們改成圖的控制流節(jié)點但有時候這種轉(zhuǎn)換會跟外部的Python狀態(tài)產(chǎn)生微妙的不一致。比如在函數(shù)里直接打印一個Tensor對象你會看到它是抽象符號而非具體值除非調(diào)用print(張量.numpy())之類的顯式操作。我也見過有人過于依賴tf.function結(jié)果把數(shù)據(jù)處理邏輯也寫進(jìn)去反而讓調(diào)試變得很麻煩。我的實踐體會是能用即時執(zhí)行模式寫清楚邏輯的時候不用急著做圖化訓(xùn)練速度不夠再針對性優(yōu)化。2.3 自動微分用GradientTape記錄梯度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的核心是反向傳播也就是求損失函數(shù)對每個參數(shù)的偏導(dǎo)數(shù)。在TensorFlow 2.x中你用tf.GradientTape這個環(huán)境管理器就能非常方便地獲得梯度。思路是把需要求導(dǎo)的計算過程包進(jìn)with tf.GradientTape() as tape:里運行時TensorFlow會記錄所有涉及tf.Variable的張量操作之后調(diào)用tape.gradient()就能根據(jù)目標(biāo)張量和變量列表算出梯度。看一個簡單的線性回歸示例import tensorflow as tf # 定義參數(shù) w tf.Variable(tf.random.normal((1,)), namew) b tf.Variable(tf.zeros((1,)), nameb) # 輸入數(shù)據(jù)這里取示例值 x tf.constant([1.0, 2.0, 3.0, 4.0]) y tf.constant([2.0, 4.0, 6.0, 8.0]) # 訓(xùn)練若干輪 for epoch in range(50): with tf.GradientTape() as tape: y_pred w * x b loss tf.reduce_mean(tf.square(y_pred - y)) grads tape.gradient(loss, [w, b]) # 手動完成一個梯度下降步驟 w.assign_sub(0.01 * grads[0]) b.assign_sub(0.01 * grads[1]) print(fw: {w.numpy()}, b: {b.numpy()})這里tape.gradient(loss, [w, b])返回的梯度列表順序和傳入變量列表完全一致。實際訓(xùn)練中優(yōu)化器比如tf.keras.optimizers.Adam的apply_gradients方法就可以直接接收這個梯度列表完成參數(shù)更新。這個機制看起來很基礎(chǔ)但它正是整條深度學(xué)習(xí)訓(xùn)練鏈路的基石。理解它之后你再去看Keras的Model.fit就不會覺得內(nèi)部有多神秘了——它其實也走的是這一套求梯度和更新參數(shù)的過程無非外面包了更多調(diào)度和狀態(tài)管理。3. 讓數(shù)據(jù)管道跑起來tf.data的高效使用心得幾乎每個用過TensorFlow的人最后都會得出同一個結(jié)論模型架構(gòu)只是深度學(xué)習(xí)系統(tǒng)的一部分如果數(shù)據(jù)管道設(shè)計得不好再好的模型也會因為喂數(shù)據(jù)太慢而在訓(xùn)練時白白等待。我自己就在這個環(huán)節(jié)踩過不少坑尤其當(dāng)數(shù)據(jù)量上了幾十萬張圖片之后每一輪訓(xùn)練里GPU有一半時間在空轉(zhuǎn)的情況非常常見。3.1 用tf.data構(gòu)建數(shù)據(jù)集的基本姿勢最基礎(chǔ)的數(shù)據(jù)集構(gòu)造方式是從內(nèi)存數(shù)組或文件路徑批量創(chuàng)建。tf.data.Dataset.from_tensor_slices常用于把圖片路徑和標(biāo)簽組合成訓(xùn)練集import tensorflow as tf data_paths [image_1.jpg, image_2.jpg, image_3.jpg] labels [0, 1, 0] dataset tf.data.Dataset.from_tensor_slices((data_paths, labels)) # 之后你可以按需對dataset做變換 dataset dataset.map(load_and_preprocess_image) dataset dataset.shuffle(1000) dataset dataset.batch(32)如果數(shù)據(jù)和標(biāo)簽都存在于內(nèi)存這種方式很直接。不過對于超大訓(xùn)練集建議先把數(shù)據(jù)寫入TFRecord格式。TFRecord是TensorFlow自定義的一種二進(jìn)制存儲格式對讀取速度友好而且可以借助tf.io.parse_single_example在數(shù)據(jù)流中進(jìn)行解碼。缺點是序列化和解析代碼寫起來有點繁瑣所以大多數(shù)小項目并不需要一上來就上TFRecord等內(nèi)存放得下的時候簡單的文件路徑列表加map變換就夠了。3.2 數(shù)據(jù)Pipeline的三大性能加速器我剛開始訓(xùn)練圖像模型時一個epoch要跑接近二十分鐘后來優(yōu)化了數(shù)據(jù)管道直接壓到了七分鐘。收益主要來自三件事并行map、預(yù)取和緩存。并行map很好理解dataset.map(load_image, num_parallel_callstf.data.AUTOTUNE)讓數(shù)據(jù)預(yù)處理解碼圖片、縮放、歸一化在多個線程里并行進(jìn)行。tf.data.AUTOTUNE是TensorFlow根據(jù)實際環(huán)境自動調(diào)整的并行度數(shù)值省去手動調(diào)參的麻煩。預(yù)取則是把下一個batch的準(zhǔn)備工作和當(dāng)前batch的計算重疊起來。調(diào)用dataset.prefetch(buffer_sizetf.data.AUTOTUNE)之后系統(tǒng)會在GPU還在訓(xùn)練當(dāng)前batch時提前準(zhǔn)備后續(xù)數(shù)據(jù)這能直接把訓(xùn)練間隙的等待時間壓到最低。我用下來的感受是prefetch是性價比最高的一個操作哪怕只是加一行延遲也能明顯下降。緩存的作用就更直接了dataset.cache()可以把第一個epoch計算過的預(yù)處理結(jié)果緩存到內(nèi)存或磁盤文件中。如果數(shù)據(jù)量不太大緩存到內(nèi)存后后續(xù)每個epoch幾乎不花預(yù)處理時間但要注意如果shuffle操作寫在cache之前緩存的就是打亂后的順序這跟你預(yù)期的隨機性可能會有出入建議把cache()放在shuffle()前面。這里我整理了一個典型的數(shù)據(jù)管道順序供你參考dataset dataset.cache() dataset dataset.shuffle(buffer_size5000) dataset dataset.map(preprocess, num_parallel_callstf.data.AUTOTUNE) dataset dataset.batch(32) dataset dataset.prefetch(tf.data.AUTOTUNE)3.3 順序問題為什么shuffle和repeat要放在恰當(dāng)位置很多人沒留意shuffle、repeat與batch之間的順序關(guān)系結(jié)果訓(xùn)練出的模型效果總差一口氣。我對原則的理解是如果希望每個epoch里都看到不同批次的組合shuffle應(yīng)該在batch之前執(zhí)行如果希望數(shù)據(jù)一直循環(huán)供應(yīng)repeat可以放在batch之前或之后取決于你想要的epoch語義。一般習(xí)慣是dataset dataset.shuffle(...).batch(...).repeat(...)這樣每個epoch內(nèi)數(shù)據(jù)是打亂的而且epoch之間數(shù)據(jù)分布也會不同。還有一個容易被忽略的問題在Dataset中如果先執(zhí)行batch再執(zhí)行shuffleshuffle的粒度就變成了隨機重排batch而不是隨機重排樣本。對于大多數(shù)監(jiān)督學(xué)習(xí)任務(wù)來說這樣會讓同一個batch里的樣本彼此高度相關(guān)訓(xùn)練初期模型的loss曲線會明顯抖動甚至影響收斂。所以我的建議是sample級別的shuffle要放在batch之前。4. 訓(xùn)練過程的可控與調(diào)優(yōu)從損失不降說起模型訓(xùn)練調(diào)試是TensorFlow使用中耗時最多的環(huán)節(jié)之一。即便代碼語法完全正確模型也常常不按預(yù)期收斂。我自己經(jīng)歷過一次印象深刻的調(diào)優(yōu)過程一個文本分類模型前五個epoch的loss一直卡在2.0左右不下去剛開始我懷疑是學(xué)習(xí)率太高導(dǎo)致震蕩于是降到1e-4結(jié)果loss紋絲不動接著懷疑模型結(jié)構(gòu)有錯又反復(fù)檢查了好幾遍仍沒發(fā)現(xiàn)問題。最后發(fā)現(xiàn)問題其實出在標(biāo)簽預(yù)處理上——我把多標(biāo)簽任務(wù)的標(biāo)簽喂成了一個形狀與模型輸出不匹配的張量損失函數(shù)的掩碼邏輯把這個不匹配悄悄掩蓋了反而導(dǎo)致梯度一直在原地打轉(zhuǎn)。4.1 損失不降的排查順序如果你也遇到loss異常平穩(wěn)或抖動不降的情況我建議按下面這個順序排查能省不少時間。先檢查數(shù)據(jù)與標(biāo)簽。打印一個batch的數(shù)據(jù)分布看看類別是否均衡標(biāo)簽是否有大量錯位歸一化是否做到位。尤其是文本序列任務(wù)padding掉的token如果直接參與損失計算也會拖住收斂這時可以借助tf.keras.layers.Masking或自定義損失函數(shù)中的mask處理。再檢查網(wǎng)絡(luò)的輸入端。輸入數(shù)據(jù)形狀與模型首層期待的shape是否一致比如卷積網(wǎng)絡(luò)輸入是不是(batch, height, width, channels)序列模型輸入是不是(batch, seq_len, features)。很多時候shape完全匹配但值域不對比如圖像沒有歸一化到[0,1]也會讓訓(xùn)練初期不穩(wěn)定。然后是優(yōu)化器與學(xué)習(xí)率。建議用tf.keras.optimizers.Adam作為默認(rèn)優(yōu)化器起步學(xué)習(xí)率先從1e-3或者3e-4這種量級試起。如果loss出現(xiàn)明顯震蕩減小學(xué)習(xí)率如果loss一直不動而不是強烈震蕩那問題多半不在學(xué)習(xí)率而在數(shù)據(jù)或梯度本身。最后可以用梯度檢查確認(rèn)模型梯度是否在正常范圍。tape.gradient得到的梯度如果出現(xiàn)大量NaN或者接近0的數(shù)值通常意味著網(wǎng)絡(luò)前面有數(shù)值不穩(wěn)定或者梯度消失。這時候圍繞網(wǎng)絡(luò)初始化、激活函數(shù)類型去排查比盲目調(diào)學(xué)習(xí)率更有用。4.2 學(xué)習(xí)率策略不是越傳越好在訓(xùn)練早期我習(xí)慣讓模型在一個較小學(xué)習(xí)率下熱身再慢慢上升到目標(biāo)學(xué)習(xí)率最后通過衰減逐步降低。TensorBoard里可以清楚看到這些策略對訓(xùn)練曲線的影響。tf.keras.optimizers.schedules提供了多種現(xiàn)成的學(xué)習(xí)率調(diào)度方式比如ExponentialDecay比較經(jīng)典CosineDecay在中大型任務(wù)上也常見。一個典型的用法是import tensorflow as tf lr_schedule tf.keras.optimizers.schedules.ExponentialDecay( initial_learning_rate1e-3, decay_steps1000, decay_rate0.96, staircaseTrue ) optimizer tf.keras.optimizers.Adam(learning_ratelr_schedule)在實際項目中我發(fā)現(xiàn)很多情況下恒定學(xué)習(xí)率加一個足夠長的預(yù)熱期效果就已經(jīng)可以了復(fù)雜的調(diào)度方案并不一定會帶來顯著提升。學(xué)習(xí)率這個超參數(shù)高度依賴數(shù)據(jù)規(guī)模和任務(wù)復(fù)雜度與其照搬別人的參數(shù)不如小規(guī)模實驗幾次然后觀察loss在訓(xùn)練過程中如何變化。4.3 過擬合問題不能靠早停一招解決很多新手一發(fā)現(xiàn)驗證集loss開始回升就立刻加上早停機制。這個做法短期有效但掩蓋了模型泛化能力不足的本質(zhì)。我更常用的一套組合拳是數(shù)據(jù)增強、權(quán)重正則化和Dropout。數(shù)據(jù)增強在圖像任務(wù)上非常好用。TensorFlow有大量內(nèi)置變換比如隨機裁剪、旋轉(zhuǎn)、翻轉(zhuǎn)、亮度對比度調(diào)整等直接在數(shù)據(jù)管道里實現(xiàn)即可。文本任務(wù)則可以用隨機的同義詞替換或弱化高頻詞等策略。權(quán)重正則化通過kernel_regularizer參數(shù)就能實現(xiàn)常見的有L1和L2正則。Dropout一般加在全連接層之后它本質(zhì)是在訓(xùn)練時隨機失活一部分神經(jīng)元迫使網(wǎng)絡(luò)學(xué)到冗余表示從而提高泛化能力。把這些手段組合起來模型即使訓(xùn)練更久也不會過早地在驗證集上惡化。5. 從訓(xùn)練到生產(chǎn)模型導(dǎo)出、量化與部署落地訓(xùn)練管道跑通只是第一步真正讓TensorFlow模型產(chǎn)生業(yè)務(wù)價值的是把模型部署到線上環(huán)境。大家在本地Jupyter里model.predict跑得順溜但一推到生產(chǎn)就各種不兼容這種情況我見得太多了。5.1 用SavedModel統(tǒng)一導(dǎo)出流程TensorFlow 2.x推薦的模型導(dǎo)出格式是SavedModel。調(diào)用model.save(my_model_dir)即可把模型的網(wǎng)絡(luò)結(jié)構(gòu)、權(quán)重與推理函數(shù)打包到一個目錄里。之后用tf.keras.models.load_model(my_model_dir)就能把模型完整恢復(fù)出來。SavedModel的好處是結(jié)構(gòu)自包含不管是用Keras定義的模型還是底層自定義的tf.Module都能統(tǒng)一導(dǎo)出而且可以同時保存多種推理簽名方便不同調(diào)用方使用。如果你的模型除了常規(guī)預(yù)測之外還需要額外的預(yù)處理或后處理邏輯可以把這些邏輯也封裝成一個函數(shù)。這里給個比較通用的服務(wù)端部署思路先用原生TensorFlow定義推理函數(shù)然后用tf.function把它和模型打包成SavedModel再借助模型服務(wù)框架對外提供HTTP或gRPC接口。import tensorflow as tf # 假設(shè)model是一個訓(xùn)練好的keras模型 model tf.keras.models.load_model(my_model_dir) tf.function(input_signature[tf.TensorSpec(shape[None], dtypetf.string)]) def serve_fn(serialized_example): parsed tf.io.parse_tensor(serialized_example, out_typetf.float32) parsed tf.reshape(parsed, [-1, 224, 224, 3]) logits model(parsed) return {probability: tf.nn.softmax(logits)} tf.saved_model.save(model, exported_model, signatures{serving_default: serve_fn})這樣部署出來的服務(wù)調(diào)用方就不需要關(guān)心模型內(nèi)部細(xì)節(jié)只需要把輸入數(shù)據(jù)按約定格式傳進(jìn)來即可。簽名機制就像給模型開了一扇指定的門不同的服務(wù)方可以訪問不同的入口。5.2 輕量化TensorFlow Lite與量化如果模型最終要跑在移動端或嵌入式設(shè)備上直接用SavedModel肯定不合適。TensorFlow Lite是面向移動端和嵌入式場景的輕量級推理框架它把模型轉(zhuǎn)換為FlatBuffer格式針對移動設(shè)備做了算子精簡和內(nèi)存優(yōu)化。轉(zhuǎn)換方式非常簡單converter tf.lite.TFLiteConverter.from_saved_model(exported_model) tflite_model converter.convert() with open(model.tflite, wb) as f: f.write(tflite_model)如果你對模型體積還比較在意可以開啟量化。量化簡單理解就是降低權(quán)重和激活值的數(shù)值精度比如從32位浮點降為8位整數(shù)換來體積減少和運行速度提升代價是精度的輕微下降。TensorFlow Lite的量化方式有動態(tài)范圍量化、完整整數(shù)量化和浮點16量化等要在部署前對比一下量化后模型與原始模型在驗證集上的指標(biāo)差多少確認(rèn)下降幅度可接受再上線。5.3 服務(wù)化部署時常見的隱藏問題有兩點我在部署階段反復(fù)踩坑特別提醒。第一訓(xùn)練階段使用tf.data處理原始數(shù)據(jù)和部署階段接收請求數(shù)據(jù)之間的預(yù)處理邏輯要保持一致。很多模型線上效果不如預(yù)期往往就是因為本地測試時走了數(shù)據(jù)管道的完整預(yù)處理而線上服務(wù)直接拿原始字符串或二進(jìn)制數(shù)據(jù)喂給模型導(dǎo)致輸入分布產(chǎn)生了偏移。解決辦法就是把預(yù)處理邏輯完整封裝到serve_fn里不管原始數(shù)據(jù)是什么形式統(tǒng)一在模型入口完成轉(zhuǎn)換。第二不同TensorFlow版本導(dǎo)出的SavedModel在更高或更低版本運行時可能產(chǎn)生算子兼容性問題。生產(chǎn)環(huán)境建議鎖定一個長期維護的TensorFlow版本更新版本前先用代表性樣本做一輪推理一致性校驗確保輸出差異在可接受范圍內(nèi)。這一點對涉及金融風(fēng)控、醫(yī)療輔助這類對數(shù)值穩(wěn)定性要求高的場景尤其重要。6. 我踩過的TensorFlow深坑與繞行方案TensorFlow作為一個生態(tài)龐大的框架雖然好用但并不意味著沒有坑。有些問題官網(wǎng)文檔里寫得并不顯眼只有自己撞上去才知道有多疼。我把印象最深的一些經(jīng)驗整理在這里算是給后來的朋友一些繞行提示。6.1 版本與Python環(huán)境的匹配問題TensorFlow對Python版本有明確的兼容范圍比如某些較老版本在Python 3.10上可能直接編譯失敗。最省心的做法是用虛擬環(huán)境單獨安裝TensorFlow并且按官方文檔給定的Python版本范圍來創(chuàng)建環(huán)境。python -m venv tf_env source tf_env/bin/activate pip install tensorflow裝好之后可以用python -c import tensorflow as tf; print(tf.__version__)確認(rèn)安裝版本。我就遇到過因為錯誤使用舊版本安裝指令導(dǎo)致裝上了CPU版TensorFlow結(jié)果GPU始終無法識別后來排查才發(fā)現(xiàn)是包名寫錯了?，F(xiàn)在GPU版本不需要單獨安裝官方主包已經(jīng)同時支持CPU和GPU只要機器環(huán)境具備CUDA和cuDNN即可。6.2 GPU顯存不釋放與分配策略訓(xùn)練一個模型之后再做推理時發(fā)現(xiàn)GPU顯存仍然被占用這是新手特別容易困惑的問題。TensorFlow默認(rèn)會預(yù)先申請盡可能多的顯存而不是按需分配。如果不設(shè)置顯存增長策略即使你的模型只需要2GB顯存TensorFlow也可能先把整個顯存吃滿。我這里說的按需可以用一行代碼配置gpus tf.config.experimental.list_physical_devices(GPU) if gpus: try: tf.config.experimental.set_memory_growth(gpus[0], True) except RuntimeError as e: print(e)set_memory_growth設(shè)置為True后顯存會隨著實際需求動態(tài)分配不會一上來就占滿。如果確實希望固定使用某塊卡的固定顯存量用tf.config.set_logical_device_configuration配合tf.config.LogicalDeviceConfiguration(memory_limit...)更靈活。6.3 多線程和進(jìn)程的程序問題TensorFlow內(nèi)部有很多線程池來控制并行算子執(zhí)行。如果你在同一個Python進(jìn)程里同時加載多個模型并頻繁新建會話容易遇到默認(rèn)線程池配置被覆蓋的問題表現(xiàn)為推理耗時突然變長。我后來的習(xí)慣是盡量復(fù)用同一個模型實例來做多次推理而不是每次預(yù)測前都重新加載一次真有多個模型需要并發(fā)處理時用多進(jìn)程而非多線程來隔離資源。另一個和線程相關(guān)的典型問題是numpy與TensorFlow混合使用時不注意張量同步點。比如在繼承tf.keras.Model的自定義call方法里如果你調(diào)用了.numpy()方法把Tensor轉(zhuǎn)成numpy數(shù)組雖然能得到數(shù)值但這個操作會打斷正在構(gòu)建的自動微分鏈路導(dǎo)致tape.gradient無法正常工作。所以在訓(xùn)練路徑里盡量別在模型內(nèi)層使用.numpy()只在外部檢查結(jié)果或做數(shù)據(jù)IO時使用。6.4 隨機種子問題帶來的不可復(fù)現(xiàn)我早年復(fù)現(xiàn)實驗結(jié)果時曾經(jīng)因為忘記設(shè)置隨機種子而白白浪費了一周時間去找模型性能波動的原因。TensorFlow的隨機性涉及多個層級Python內(nèi)置的random模塊、numpy的RandomState、TensorFlow自身的圖隨機算子。想要最大程度地固定隨機行為需要三層同時設(shè)import os import random import numpy as np import tensorflow as tf os.environ[PYTHONHASHSEED] 0 random.seed(42) np.random.seed(42) tf.random.set_seed(42)不過要說明的是即使都設(shè)了種子使用GPU并行算子時某些情況下仍然可能存在細(xì)微的非確定性輸出。如果科研或工程上對可復(fù)現(xiàn)性要求很高可以考慮限制部分算子使用CPU執(zhí)行或者在代碼里記錄關(guān)鍵變量的hash值方便對照。寫在最后的幾句真心話如果讓我給剛接觸TensorFlow的人提一個最實用的建議那一定是先沉住氣把張量形狀、dtype和各種常見的張量操作練熟再去看模型結(jié)構(gòu)。TensorFlow的語法體系非常龐雜但絕大多數(shù)操作歸根結(jié)底都是圍繞張量展開的。你會發(fā)現(xiàn)一旦具備了這種形狀思維讀官方文檔就像查字典看到某個API名字就能大致猜到它的輸入輸出形態(tài)。我自己在實際操作中還養(yǎng)成了一個習(xí)慣任何新方案上線前先用小數(shù)據(jù)集跑通一個最小實驗再逐步放大。這個習(xí)慣幫我在不少項目中及時發(fā)現(xiàn)了版本兼容和數(shù)據(jù)處理上的隱藏問題成本遠(yuǎn)比直接跑全量訓(xùn)練低得多。TensorFlow的生態(tài)仍在持續(xù)演進(jìn)新的API和最佳實踐層出不窮但核心思想是穩(wěn)定的。把這個框架用好靠的從來不是刷多少教程而是在一次次訓(xùn)練、調(diào)優(yōu)和部署的過程中積累起來的手感。