配置與梯度沖突實(shí)戰(zhàn)指南)
1. 多輸出模型里損失函數(shù)到底該怎么擺很多人第一次在Keras里給模型掛多個(gè)損失函數(shù)腦子里想的都是“我只要把loss寫成一個(gè)列表Keras自己會(huì)搞定”。結(jié)果跑起來要么報(bào)形狀不匹配要么某個(gè)分支的loss壓根沒參與梯度更新要么訓(xùn)練半天發(fā)現(xiàn)只有主任務(wù)在收斂輔助任務(wù)紋絲不動(dòng)。問題不在于Keras不支持多損失而在于大多數(shù)人沒搞清楚Keras處理多損失時(shí)的三條路徑輸出層命名、loss字典映射、loss_weights權(quán)重分配。這三者只要有一個(gè)對(duì)不上整個(gè)訓(xùn)練圖就會(huì)悄悄退化。先明確一個(gè)概念Keras里的“多個(gè)自定義損失函數(shù)”其實(shí)分兩種場(chǎng)景。第一種是單輸出多損失比如一個(gè)輸出既要算MSE又要算一個(gè)正則項(xiàng)這時(shí)候你需要的是把多個(gè)loss加起來返回一個(gè)標(biāo)量。第二種是多輸出多損失模型有多個(gè)head每個(gè)head配一個(gè)獨(dú)立的loss這時(shí)候才涉及l(fā)oss列表或字典。很多人把這兩種場(chǎng)景混在一起寫代碼自然亂。我見過最典型的錯(cuò)誤寫法是這樣的model.compile(optimizeradam, loss[my_loss_a, my_loss_b], loss_weights[0.7, 0.3])但模型只有一個(gè)輸出層。這種情況下Keras會(huì)嘗試把兩個(gè)loss都作用在同一個(gè)輸出上如果兩個(gè)loss返回的形狀不一致直接報(bào)維度錯(cuò)誤如果形狀一致它會(huì)做加權(quán)求和但你以為的“兩個(gè)獨(dú)立任務(wù)”其實(shí)變成了“同一個(gè)預(yù)測(cè)值的兩個(gè)約束”語義完全變了。正確的判斷邏輯是先看model.outputs的長度。如果長度是1那所有l(wèi)oss最終必須聚合成一個(gè)標(biāo)量如果長度大于1才用列表或字典給每個(gè)輸出分配loss。這個(gè)判斷我建議你在寫代碼前先用print(len(model.outputs))確認(rèn)一遍別憑感覺。還有一個(gè)隱藏坑當(dāng)模型有多個(gè)輸出但只給部分輸出配了lossKeras默認(rèn)會(huì)對(duì)沒有配loss的輸出不做梯度回傳。這本身是合理的但如果你用的是函數(shù)式API且某個(gè)中間層被多個(gè)輸出共享沒配loss的那個(gè)分支仍然會(huì)參與前向計(jì)算只是不產(chǎn)生梯度。這會(huì)導(dǎo)致顯存占用比預(yù)期高訓(xùn)練速度也上不去。解決辦法是顯式給每個(gè)輸出都配loss哪怕某個(gè)分支你暫時(shí)不想訓(xùn)練也給它配一個(gè)權(quán)重為0的loss這樣Keras的圖結(jié)構(gòu)會(huì)更干凈。提示多輸出模型里輸出層的name參數(shù)非常關(guān)鍵。后面用字典配loss時(shí)key必須和輸出層的name完全一致大小寫敏感。2. 自定義損失函數(shù)的三種寫法與適用邊界Keras里寫自定義損失函數(shù)表面上有三種寫法純函數(shù)式、帶額外參數(shù)的閉包式、繼承Loss類。很多人隨手選一種就開始寫結(jié)果在保存模型、加載模型、分布式訓(xùn)練時(shí)接連翻車。這三種寫法不是等價(jià)的它們的序列化能力、參數(shù)綁定方式、對(duì)add_loss的兼容性都不一樣。2.1 純函數(shù)式最簡單但最受限純函數(shù)式就是接收y_true和y_pred兩個(gè)參數(shù)返回一個(gè)標(biāo)量或逐樣本的loss張量import tensorflow as tf def focal_loss(y_true, y_pred, gamma2.0, alpha0.25): y_pred tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7) ce -y_true * tf.math.log(y_pred) weight alpha * tf.pow(1 - y_pred, gamma) return tf.reduce_mean(weight * ce)這種寫法的問題在于gamma和alpha被硬編碼在函數(shù)簽名里編譯時(shí)沒法從外部傳參。如果你在model.compile(lossfocal_loss)時(shí)想改gamma只能重新定義一個(gè)新函數(shù)。更麻煩的是模型保存后load_model需要能定位到這個(gè)函數(shù)如果函數(shù)定義在__main__里且沒有注冊(cè)加載時(shí)會(huì)報(bào)Unknown loss function。適用邊界超參數(shù)固定、不需要頻繁調(diào)整、且函數(shù)定義在可導(dǎo)入的模塊文件里。如果你只是在notebook里快速實(shí)驗(yàn)純函數(shù)式夠用但一旦要交付或部署建議換寫法。2.2 閉包式參數(shù)靈活但序列化有坑閉包式是用一個(gè)外層函數(shù)包住超參數(shù)返回內(nèi)層loss函數(shù)def make_focal_loss(gamma2.0, alpha0.25): def focal_loss(y_true, y_pred): y_pred tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7) ce -y_true * tf.math.log(y_pred) weight alpha * tf.pow(1 - y_pred, gamma) return tf.reduce_mean(weight * ce) return focal_loss這樣你可以lossmake_focal_loss(gamma3.0)參數(shù)靈活。但保存模型時(shí)Keras序列化的是內(nèi)層函數(shù)的引用而內(nèi)層函數(shù)是運(yùn)行時(shí)動(dòng)態(tài)創(chuàng)建的load_model根本找不到它。你會(huì)遇到經(jīng)典的ValueError: Unknown loss function: focal_loss。解決辦法有兩個(gè)一是用custom_objects在加載時(shí)手動(dòng)傳入二是改用Loss類。前者在簡單場(chǎng)景可行但如果你有多個(gè)閉包losscustom_objects的字典會(huì)變得很難維護(hù)。2.3 繼承Loss類最規(guī)范但代碼量最大tf.keras.utils.register_keras_serializable(packageMyLosses) class FocalLoss(tf.keras.losses.Loss): def __init__(self, gamma2.0, alpha0.25, **kwargs): super().__init__(**kwargs) self.gamma gamma self.alpha alpha def call(self, y_true, y_pred): y_pred tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7) ce -y_true * tf.math.log(y_pred) weight self.alpha * tf.pow(1 - y_pred, self.gamma) return tf.reduce_mean(weight * ce) def get_config(self): config super().get_config() config.update({gamma: self.gamma, alpha: self.alpha}) return config這種寫法配合register_keras_serializable裝飾器保存和加載都不會(huì)出問題。get_config負(fù)責(zé)把超參數(shù)寫進(jìn)模型文件加載時(shí)自動(dòng)重建。代價(jià)是代碼量翻倍而且call方法里不能直接用Python的if判斷張量值必須用tf.cond或tf.where。我的建議是實(shí)驗(yàn)階段用閉包式交付階段一律換成Loss類。切換成本不高但能省掉后面一堆序列化相關(guān)的調(diào)試時(shí)間。寫法參數(shù)靈活性序列化支持代碼量推薦場(chǎng)景純函數(shù)式低需手動(dòng)注冊(cè)最少快速驗(yàn)證閉包式高差中等實(shí)驗(yàn)調(diào)參Loss類高原生支持最多生產(chǎn)交付3. 多輸出模型的loss字典配置與權(quán)重調(diào)優(yōu)當(dāng)模型有多個(gè)輸出時(shí)最清晰的配置方式是用字典把輸出層name映射到對(duì)應(yīng)的lossmodel.compile( optimizertf.keras.optimizers.Adam(1e-3), loss{ classification_head: FocalLoss(gamma2.0), regression_head: tf.keras.losses.Huber(delta1.0), auxiliary_head: tf.keras.losses.MeanSquaredError() }, loss_weights{ classification_head: 1.0, regression_head: 0.5, auxiliary_head: 0.2 }, metrics{ classification_head: [accuracy], regression_head: [mae] } )這里有幾個(gè)細(xì)節(jié)值得展開。第一loss_weights的key必須和loss字典的key完全一致否則Keras會(huì)忽略權(quán)重默認(rèn)全部按1.0處理。第二metrics也是按輸出層name配置的如果你只給部分輸出配了metrics訓(xùn)練日志里只會(huì)顯示配了的那些沒配的不會(huì)報(bào)錯(cuò)但也不會(huì)顯示。第三loss_weights的總和不需要等于1Keras內(nèi)部會(huì)做歸一化處理但歸一化方式是把每個(gè)loss乘以對(duì)應(yīng)權(quán)重后直接相加不做除法。所以如果你把所有權(quán)重都放大10倍等效于把學(xué)習(xí)率放大10倍這一點(diǎn)很容易被忽略。權(quán)重怎么調(diào)我的經(jīng)驗(yàn)是分三步走。第一步先把所有l(wèi)oss_weights設(shè)為1.0跑幾個(gè)epoch觀察每個(gè)loss的初始量級(jí)。如果分類loss在2.0左右回歸loss在0.01左右那回歸任務(wù)的梯度貢獻(xiàn)幾乎可以忽略。第二步根據(jù)量級(jí)差異反比設(shè)置權(quán)重讓每個(gè)loss的初始貢獻(xiàn)大致相當(dāng)。比如分類loss是2.0回歸loss是0.01那回歸權(quán)重可以設(shè)成100左右但不要一步到位先設(shè)成10觀察。第三步根據(jù)任務(wù)重要性微調(diào)。如果分類是主任務(wù)回歸是輔助那在量級(jí)對(duì)齊的基礎(chǔ)上把分類權(quán)重再放大1.5到2倍。注意loss_weights調(diào)得過大時(shí)某個(gè)分支的梯度會(huì)主導(dǎo)整個(gè)網(wǎng)絡(luò)導(dǎo)致共享層被這個(gè)任務(wù)“綁架”其他任務(wù)的表現(xiàn)反而下降。這種情況在多任務(wù)學(xué)習(xí)里叫“負(fù)遷移”表現(xiàn)是某個(gè)任務(wù)的驗(yàn)證指標(biāo)先降后升再降。遇到這種情況要么降低該任務(wù)權(quán)重要么在共享層和任務(wù)頭之間加梯度裁剪。還有一個(gè)實(shí)操技巧用tf.keras.callbacks.Callback在訓(xùn)練過程中動(dòng)態(tài)打印每個(gè)loss的值而不是只看總loss??俵oss是加權(quán)和某個(gè)分支出問題時(shí)總loss可能看起來正常但分支loss已經(jīng)發(fā)散。下面這個(gè)callback我經(jīng)常用class LossLogger(tf.keras.callbacks.Callback): def on_epoch_end(self, epoch, logsNone): logs logs or {} for key, value in logs.items(): if loss in key and val not in key: print(fEpoch {epoch} | {key}: {value:.4f})這樣你能清楚看到每個(gè)head的loss走勢(shì)比只看一個(gè)total_loss有用得多。4. 梯度沖突與loss平衡的實(shí)戰(zhàn)處理多損失訓(xùn)練最棘手的問題不是配置而是梯度沖突。兩個(gè)任務(wù)的loss對(duì)共享層參數(shù)的梯度方向相反時(shí)網(wǎng)絡(luò)會(huì)在兩個(gè)任務(wù)之間反復(fù)橫跳表現(xiàn)是總loss震蕩不降或者一個(gè)任務(wù)好另一個(gè)任務(wù)差。這個(gè)問題在推薦系統(tǒng)、多模態(tài)模型里特別常見。4.1 用GradNorm思路做動(dòng)態(tài)權(quán)重固定權(quán)重最大的問題是訓(xùn)練過程中各任務(wù)的學(xué)習(xí)難度在變化前期簡單的任務(wù)梯度大后期困難的任務(wù)梯度大固定權(quán)重沒法適應(yīng)。GradNorm的核心思路是根據(jù)每個(gè)任務(wù)當(dāng)前的梯度范數(shù)和訓(xùn)練速度動(dòng)態(tài)調(diào)整loss權(quán)重讓各任務(wù)的梯度范數(shù)保持在同一量級(jí)。簡化版的實(shí)現(xiàn)思路是在每個(gè)batch結(jié)束后計(jì)算每個(gè)任務(wù)loss對(duì)共享層參數(shù)的梯度范數(shù)然后按范數(shù)的反比更新權(quán)重。完整實(shí)現(xiàn)比較復(fù)雜但你可以用一個(gè)近似版本class DynamicWeightCallback(tf.keras.callbacks.Callback): def __init__(self, shared_layer_name, task_names, alpha0.1): super().__init__() self.shared_layer_name shared_layer_name self.task_names task_names self.alpha alpha self.weights {name: 1.0 for name in task_names} def on_train_batch_end(self, batch, logsNone): logs logs or {} losses {} for name in self.task_names: key f{name}_loss if key in logs: losses[name] logs[key] if len(losses) 2: return mean_loss sum(losses.values()) / len(losses) for name, loss_val in losses.items(): ratio loss_val / (mean_loss 1e-8) self.weights[name] * (1 - self.alpha * (ratio - 1)) self.weights[name] max(0.1, min(10.0, self.weights[name]))這個(gè)callback在on_train_batch_end里根據(jù)各任務(wù)loss與平均loss的比值調(diào)整權(quán)重loss高于平均的任務(wù)權(quán)重降低低于平均的權(quán)重升高。alpha控制調(diào)整速度太大會(huì)震蕩太小沒效果一般從0.05到0.1開始試。權(quán)重上下限是為了防止某個(gè)任務(wù)權(quán)重被壓到0或漲到無窮。4.2 梯度裁剪與共享層保護(hù)如果動(dòng)態(tài)權(quán)重調(diào)起來太麻煩一個(gè)更直接的辦法是在共享層后面加梯度裁剪。Keras里可以在optimizer層面做全局裁剪optimizer tf.keras.optimizers.Adam( learning_rate1e-3, clipnorm1.0 )clipnorm按梯度向量的L2范數(shù)裁剪clipvalue按每個(gè)梯度元素的值裁剪。多任務(wù)場(chǎng)景下我更推薦clipnorm因?yàn)樗A袅颂荻确较蛑幌拗品?。clipvalue會(huì)改變梯度方向可能導(dǎo)致某些任務(wù)學(xué)不動(dòng)。另一個(gè)技巧是給共享層加一個(gè)較小的學(xué)習(xí)率給任務(wù)頭加較大的學(xué)習(xí)率。Keras原生不支持分層學(xué)習(xí)率但可以用tf.keras.optimizers.schedules配合自定義訓(xùn)練循環(huán)實(shí)現(xiàn)。簡單做法是共享層用SGD任務(wù)頭用Adam但這樣需要寫自定義train_step。如果不想寫自定義循環(huán)退而求其次的辦法是在共享層后面加BatchNormalizationBN層的running statistics會(huì)起到一定的穩(wěn)定作用緩解梯度沖突帶來的震蕩。4.3 任務(wù)優(yōu)先級(jí)與課程學(xué)習(xí)有時(shí)候梯度沖突的根源是任務(wù)太難模型在早期根本學(xué)不動(dòng)。這時(shí)候可以用課程學(xué)習(xí)的思路前期只訓(xùn)練主任務(wù)等主任務(wù)收斂到一定程度后再加入輔助任務(wù)。實(shí)現(xiàn)方式是用callback在指定epoch后修改model.loss_weightsclass CurriculumCallback(tf.keras.callbacks.Callback): def __init__(self, switch_epoch, new_weights): super().__init__() self.switch_epoch switch_epoch self.new_weights new_weights def on_epoch_begin(self, epoch, logsNone): if epoch self.switch_epoch: self.model.loss_weights self.new_weights print(fEpoch {epoch}: loss weights switched to {self.new_weights})注意model.loss_weights在編譯后修改是否生效取決于TF版本較新的版本支持動(dòng)態(tài)修改老版本可能需要重新編譯。實(shí)測(cè)下來TF 2.8以上都可以直接改但改完后第一個(gè)batch的loss可能會(huì)跳變屬于正?,F(xiàn)象。5. 保存加載與分布式訓(xùn)練的兼容性細(xì)節(jié)自定義loss在單機(jī)單卡上跑通不難難的是保存加載和分布式訓(xùn)練時(shí)不掉鏈子。我踩過的坑里80%都出在這兩個(gè)環(huán)節(jié)。5.1 保存加載的三種失敗模式第一種是Unknown loss function原因是loss函數(shù)沒有注冊(cè)。解決辦法是用tf.keras.utils.register_keras_serializable裝飾Loss類或者在load_model時(shí)傳custom_objects。我強(qiáng)烈推薦前者因?yàn)楹笳咴谀P捅黄渌a加載時(shí)容易遺漏。第二種是get_config不完整。如果你在Loss類里加了自定義屬性但忘了寫進(jìn)get_config保存時(shí)不會(huì)報(bào)錯(cuò)加載后屬性會(huì)丟失loss計(jì)算結(jié)果和訓(xùn)練時(shí)不一致。這種bug最隱蔽因?yàn)槟P湍芗虞d、能推理只是結(jié)果悄悄偏了。檢查方法是保存后重新加載用同一批數(shù)據(jù)分別跑原模型和加載模型對(duì)比loss值是否完全一致。第三種是loss_weights丟失。Keras保存模型時(shí)會(huì)把compile信息寫進(jìn)文件但如果你在保存后手動(dòng)改了loss_weights再保存舊值可能被覆蓋。穩(wěn)妥做法是每次保存前重新compile一次確保配置是最新的。# 保存前重新compile確保配置寫入 model.compile(optimizeroptimizer, lossloss_dict, loss_weightsweights) model.save(my_model.keras) # 加載時(shí)如果用了register_keras_serializable不需要custom_objects loaded tf.keras.models.load_model(my_model.keras)5.2 分布式訓(xùn)練下的loss聚合用MirroredStrategy做多卡訓(xùn)練時(shí)每個(gè)卡上的loss是各自計(jì)算的Keras會(huì)自動(dòng)做all-reduce求平均。但如果你在loss函數(shù)里用了tf.reduce_sum而不是tf.reduce_mean多卡下的總loss會(huì)被放大N倍N是卡數(shù)等效于學(xué)習(xí)率放大N倍。這個(gè)坑我在一次四卡訓(xùn)練里踩過單卡正常四卡直接發(fā)散。解決辦法很簡單loss函數(shù)里一律用tf.reduce_mean讓Keras的分布式聚合邏輯處理跨卡平均。如果你確實(shí)需要sum那就在compile時(shí)把學(xué)習(xí)率除以卡數(shù)。另一個(gè)分布式相關(guān)的細(xì)節(jié)是loss_weights在MirroredStrategy下的行為。權(quán)重是在聚合之后應(yīng)用的所以每張卡上的加權(quán)方式一致不需要額外處理。但如果你用了自定義訓(xùn)練循環(huán)需要手動(dòng)在strategy.run外面做加權(quán)否則每張卡會(huì)各自加權(quán)再聚合結(jié)果和單卡不一致。提示分布式訓(xùn)練時(shí)建議先用單卡跑通完整流程確認(rèn)loss收斂后再切多卡。直接上多卡調(diào)試問題定位難度會(huì)翻倍。6. 幾個(gè)我實(shí)際踩過的坑和對(duì)應(yīng)解法第一個(gè)坑在loss函數(shù)里用了Python的if判斷張量值。比如if y_pred 0.5: return ...這在eager模式下能跑但一旦用tf.function編譯或保存成SavedModel直接報(bào)OperatorNotAllowedInGraphError。正確做法是用tf.where# 錯(cuò)誤寫法 def bad_loss(y_true, y_pred): if tf.reduce_mean(y_pred) 0.5: return tf.reduce_mean(tf.square(y_true - y_pred)) return tf.reduce_mean(tf.abs(y_true - y_pred)) # 正確寫法 def good_loss(y_true, y_pred): mse tf.reduce_mean(tf.square(y_true - y_pred)) mae tf.reduce_mean(tf.abs(y_true - y_pred)) condition tf.reduce_mean(y_pred) 0.5 return tf.where(condition, mse, mae)第二個(gè)坑loss函數(shù)返回了逐樣本的loss向量而不是標(biāo)量。Keras的compile接口期望loss返回標(biāo)量如果你返回形狀為(batch_size,)的張量Keras會(huì)做一次隱式reduce_mean但如果你在loss里已經(jīng)做了加權(quán)這個(gè)隱式平均會(huì)打亂權(quán)重。表現(xiàn)是訓(xùn)練能跑但收斂異常。檢查方法是打印loss函數(shù)的返回形狀確保是()而不是(batch_size,)。第三個(gè)坑多個(gè)loss共享同一個(gè)可訓(xùn)練變量時(shí)梯度累加導(dǎo)致更新量翻倍。比如兩個(gè)loss都對(duì)同一個(gè)embedding層求梯度Keras會(huì)把兩個(gè)梯度相加后更新。這本身是合理的但如果你期望的是兩個(gè)loss分別更新就需要用tf.GradientTape手動(dòng)控制。大多數(shù)情況下梯度相加是想要的行為但如果你發(fā)現(xiàn)某個(gè)共享層更新過快可以檢查是不是多個(gè)loss都在對(duì)它產(chǎn)生梯度。第四個(gè)坑metrics和loss用了同一個(gè)函數(shù)但期望不同行為。比如你把focal loss同時(shí)配成loss和metricmetric會(huì)顯示focal loss的值但這個(gè)值受loss_weights影響而metric本身不應(yīng)該被加權(quán)。結(jié)果是metric數(shù)值和實(shí)際任務(wù)表現(xiàn)對(duì)不上。正確做法是loss和metric分開定義metric用標(biāo)準(zhǔn)的accuracy或mae??颖憩F(xiàn)根因解法Python if判斷張量Graph模式報(bào)錯(cuò)圖模式不支持Python控制流改用tf.whereloss返回向量收斂異常隱式reduce打亂權(quán)重確保返回標(biāo)量共享層梯度翻倍更新過快多l(xiāng)oss梯度累加檢查是否預(yù)期行為metric被加權(quán)指標(biāo)失真metric受loss_weights影響loss和metric分開7. 一套可直接復(fù)用的多損失模板把上面所有經(jīng)驗(yàn)整合成一個(gè)模板你可以直接抄去改。這個(gè)模板覆蓋了Loss類定義、多輸出配置、動(dòng)態(tài)權(quán)重、保存加載四個(gè)環(huán)節(jié)。import tensorflow as tf tf.keras.utils.register_keras_serializable(packageMyLosses) class WeightedFocalLoss(tf.keras.losses.Loss): def __init__(self, gamma2.0, alpha0.25, nameweighted_focal_loss, **kwargs): super().__init__(namename, **kwargs) self.gamma gamma self.alpha alpha def call(self, y_true, y_pred): y_pred tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7) ce -y_true * tf.math.log(y_pred) weight self.alpha * tf.pow(1 - y_pred, self.gamma) return tf.reduce_mean(weight * ce) def get_config(self): config super().get_config() config.update({gamma: self.gamma, alpha: self.alpha}) return config def build_model(): inputs tf.keras.Input(shape(128,), nameinput) shared tf.keras.layers.Dense(256, activationrelu, nameshared)(inputs) shared tf.keras.layers.BatchNormalization(nameshared_bn)(shared) shared tf.keras.layers.Dropout(0.3)(shared) cls_out tf.keras.layers.Dense(10, activationsoftmax, namecls_head)(shared) reg_out tf.keras.layers.Dense(1, namereg_head)(shared) return tf.keras.Model(inputsinputs, outputs[cls_out, reg_out]) model build_model() loss_dict { cls_head: WeightedFocalLoss(gamma2.0, alpha0.25), reg_head: tf.keras.losses.Huber(delta1.0) } loss_weights { cls_head: 1.0, reg_head: 0.5 } model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3, clipnorm1.0), lossloss_dict, loss_weightsloss_weights, metrics{cls_head: [accuracy], reg_head: [mae]} ) model.summary() # 保存 model.save(multi_loss_model.keras) # 加載因?yàn)橛昧藃egister_keras_serializable不需要custom_objects loaded_model tf.keras.models.load_model(multi_loss_model.keras)這個(gè)模板里幾個(gè)關(guān)鍵點(diǎn)再強(qiáng)調(diào)一遍clipnorm1.0防止梯度爆炸BatchNormalization穩(wěn)定共享層register_keras_serializable保證序列化loss_weights按任務(wù)重要性分配。你拿到后只需要改輸入形狀、輸出維度、loss類型和權(quán)重值。如果訓(xùn)練過程中發(fā)現(xiàn)某個(gè)head的loss不降先檢查三件事該head的輸出層name是否和loss字典的key一致、該head的loss是否返回標(biāo)量、該head的梯度是否被其他head的梯度淹沒。這三件事查完90%的問題都能定位。最后分享一個(gè)我常用的調(diào)試技巧在訓(xùn)練前用model.compute_loss手動(dòng)算一次loss傳入隨機(jī)生成的y_true和y_pred看返回的loss字典結(jié)構(gòu)是否符合預(yù)期。這一步能在正式訓(xùn)練前暴露大部分配置錯(cuò)誤比跑幾個(gè)epoch再發(fā)現(xiàn)問題省時(shí)間得多。