：用Keras實現(xiàn)穩(wěn)定GAN訓練的梯度懲罰方案)
如果你也經(jīng)歷過GAN訓練中那種玄學式的崩潰——生成器輸出一堆噪點、loss曲線像心電圖一樣亂跳、換一個隨機種子結(jié)果就天差地別——那么WGAN-GP和WGAN-Div這兩套方案應(yīng)該能幫你從泥潭里爬出來。我最近用Keras在MNIST和CIFAR-10上把這兩個模型完整搭了一遍都能穩(wěn)定生成圖像。這篇文章就把從原理到代碼再到調(diào)參的整個過程記錄下來給想動手復現(xiàn)的人一個可以直接抄作業(yè)的參考。先說明一個前提下面所有代碼基于TensorFlow 2.x內(nèi)置的Keras。如果你還在用獨立的Keras 2.xAPI基本通用但GradientTape那一套只存在于tf.keras里建議直接切換到TensorFlow 2.x省掉版本兼容的折騰。1. 為什么繞不開WGAN這條路線1.1 經(jīng)典GAN的病灶JS散度在偷懶傳統(tǒng)GAN里判別器輸出的是一個概率用交叉熵做loss本質(zhì)上是在最小化生成分布和真實分布之間的JS散度。這個度量方式有個致命問題當兩個分布的重疊區(qū)域非常小的時候JS散度會趨于一個常數(shù)log2梯度幾乎為0。舉個例子。假設(shè)真實數(shù)據(jù)分布是二維平面上的一條細線生成器剛開始只能撒出一些離散的噪點。這兩組點在絕大多數(shù)位置都不重合判別器很容易就把它們分得明明白白。此時生成器幾乎收不到任何有用的梯度信號不管怎么調(diào)參數(shù)輸出的還是那堆噪點。這就是大家常說的“判別器太強生成器被錘死”。實際操作中的表現(xiàn)就是判別器的loss快速降到接近0生成器的loss卻不怎么動生成的圖片越來越像噪聲。很多人第一反應(yīng)是調(diào)判別器的學習率但根子在于度量方式本身就給不出有效的“距離感”。1.2 Wasserstein距離換了個度量方式WGAN的核心改動是把判別器換成了“critic”輸出不再是一個概率而是一個實數(shù)得分。它優(yōu)化的目標也從JS散度換成Wasserstein距離也就是常說的“推土機距離”。這個距離的直覺很好理解假設(shè)真實分布是一堆土生成分布是另一堆土Wasserstein距離就是把這堆土搬到另一堆土的最小運輸成本。即使兩堆土完全不重疊只要它們位置不同這個成本就是一個正的有限值梯度就不會消失。這就解決了經(jīng)典GAN里“距離太遠就無梯度”的問題。所以WGAN的critic要做的事情不是判斷“真”“假”而是給真實樣本打高分給生成樣本打低分。對應(yīng)到代碼里就是去掉最后一層的sigmoid輸出一個raw score。這是WGAN與傳統(tǒng)GAN最直觀的區(qū)別。1.3 權(quán)重裁剪的副作用WGAN為了保證Wasserstein距離能算得下去要求critic的函數(shù)滿足1-Lipschitz條件也就是輸出對輸入的改變不能過于敏感。論文作者最初用的方法是權(quán)重裁剪把critic的每一層權(quán)重強行限制在一個小范圍內(nèi)比如[-0.01, 0.01]。這個方法實現(xiàn)起來非常簡單但副作用很明顯。權(quán)重被強行裁剪后critic的表達能力被壓縮得很厲害很多情況下它只能學到非常粗糙的判別函數(shù)。我一開始照著原始WGAN代碼跑發(fā)現(xiàn)生成的圖像總有一種“模糊感”細節(jié)上不去。后來才知道這是權(quán)重裁剪帶來的典型問題critic太弱給生成器提供的梯度質(zhì)量也差。WGAN-GP和WGAN-Div就是在這個基礎(chǔ)上提出的兩種改進。它們的核心思路一模一樣不再裁剪權(quán)重而是通過給損失函數(shù)加一個“梯度懲罰項”直接約束critic對輸入的梯度從而滿足Lipschitz條件。2. 實驗整體設(shè)計先定結(jié)構(gòu)再談?chuàng)p失2.1 數(shù)據(jù)準備與歸一化我用兩個數(shù)據(jù)集做驗證MNIST用來快速驗證收斂邏輯CIFAR-10用來驗證彩色、復雜圖像上的生成效果。兩個數(shù)據(jù)集都縮放到[-1, 1]區(qū)間生成器最后一層用tanh激活輸出自然落在[-1, 1]內(nèi)正好對應(yīng)輸入的數(shù)據(jù)范圍。這里有個容易忽略的細節(jié)如果你把圖像歸一化到[0, 1]生成器最后一層還在用tanh輸出那生成圖像的像素會被強制壓到[-1, 1]看起來會“發(fā)灰”。反過來如果數(shù)據(jù)在[-1, 1]內(nèi)生成器卻用sigmoid輸出那圖像永遠達不到最大亮度和對比度。所以訓練前先確認好數(shù)據(jù)的范圍與輸出層的激活函數(shù)。2.2 生成器結(jié)構(gòu)生成器沿用DCGAN那套結(jié)構(gòu)先通過全連接層把隨機噪聲映射成一個小尺寸的特征圖再用轉(zhuǎn)置卷積逐步上采樣到目標尺寸。我在CIFAR-10上用32×32分辨率從8×8的特征圖開始連續(xù)兩次stride為2的轉(zhuǎn)置卷積到16×16、32×32最后一層改用普通卷積輸出3個通道。def build_generator(latent_dim128): noise Input(shape(latent_dim,)) x Dense(8 * 8 * 256, use_biasFalse)(noise) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Reshape((8, 8, 256))(x) x Conv2DTranspose(128, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2DTranspose(64, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2D(3, (3, 3), paddingsame, activationtanh)(x) return Model(noise, x)生成器里每一層都用LeakyReLUalpha取0.2。這個選擇是為了防止ReLU死掉——GAN訓練過程中生成器的輸入分布變化很大ReLU一旦被壓到負區(qū)間就可能一直輸出0導致梯度斷流。LeakyReLU可以保證梯度至少有一條通道能流回去。BN層在生成器里很重要能顯著穩(wěn)定訓練。但注意調(diào)用生成器時如果處于訓練階段一定要寫generator(noise, trainingTrue)讓BN更新批次統(tǒng)計量。生成圖像做推理時再用trainingFalse否則你會發(fā)現(xiàn)同一個噪聲向量每次生成的結(jié)果不一樣原因是BN在訓練模式下用了當前batch的均值方差。2.3 Critic結(jié)構(gòu)critic的結(jié)構(gòu)與生成器鏡像但有幾處關(guān)鍵差異不使用BN不使用sigmoid最后一層直接輸出一個標量。def build_critic(img_shape(32, 32, 3)): img Input(shapeimg_shape) x Conv2D(64, (4, 4), strides(2, 2), paddingsame)(img) x LeakyReLU(alpha0.2)(x) x Conv2D(128, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Conv2D(256, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Flatten()(x) x Dropout(0.2)(x) out Dense(1)(x) return Model(img, out)為什么critic不用BN梯度懲罰的計算要求在插值樣本上對輸入求導而BN在訓練時引入批內(nèi)統(tǒng)計量使輸出不再只依賴當前單個輸入相當于在批維度上產(chǎn)生了信息泄露。這樣一來“對單個輸入的梯度”就變得不再純粹懲罰項也會受到污染。很多第一次復現(xiàn)WGAN-GP的人把BN留在critic里結(jié)果訓練不穩(wěn)定多半是這個原因。如果卷積下采樣想保留更多信息也可以不加Dropout。我加了一層Dropout(0.2)是防止critic對訓練集過度敏感實測對穩(wěn)定性有一點幫助但不是必須的。2.4 訓練超參整套實驗的關(guān)鍵超參如下latent_dim 128batch_size 64critic更新次數(shù)n_critic 5MNIST上我試過1也能收斂優(yōu)化器Adamlr 1e-4beta_1 0.5WGAN-GP的lambda 10WGAN-Div的k 1p 2起步Adam的beta_1取0.5而不是默認的0.9這里說下原因。GAN的loss曲面本身波動劇烈如果動量累積過大梯度更新方向會“剎不住車”導致訓練震蕩。beta_1降低以后動量衰減更快更新方向更追隨當前梯度穩(wěn)定性會好很多。n_critic是一個可以調(diào)的空閑參數(shù)。理論上critic需要先訓練充分才能給生成器提供準確的Wasserstein梯度。但如果critic太強、生成器太弱同樣會出現(xiàn)生成器完全跟不上的問題。我在小數(shù)據(jù)集上的經(jīng)驗是n_critic1到5都能跑數(shù)據(jù)集越小、越簡單n_critic就可以越小。3. WGAN-GP梯度懲罰的完整實現(xiàn)3.1 插值樣本與梯度懲罰的計算WGAN-GP的核心是構(gòu)造一個“插值樣本”讓critic在這個樣本上的梯度范數(shù)盡量接近1。所謂插值樣本就是從一對真實樣本和生成樣本的連線上隨機取一個點alpha * real (1 - alpha) * fakealpha是[0,1]間的隨機數(shù)。為什么要在連線上采樣因為Lipschitz條件要求函數(shù)的梯度在全空間都有界但真正計算所有位置不可行。經(jīng)驗上真實樣本和生成樣本中間過渡區(qū)域是最容易違反Lipschitz條件的地方在連線上做懲罰性價比最高。這個思想不是嚴格的數(shù)學證明但實驗效果非常好。def gradient_penalty(critic, real, fake): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3])) gp tf.reduce_mean((norm - 1.0) ** 2) return gp這里有個初學者容易出錯的地方tf.GradientTape()里除了要計算critic的輸出還要顯式調(diào)用tape.watch(interp)否則TensorFlow不會自動記錄對interp的梯度。很多報錯“gradient is None”都是因為這個。3.2 損失函數(shù)與訓練循環(huán)WGAN-GP中critic的損失不是交叉熵而是三部分之和對真實樣本的輸出取平均對生成樣本的輸出取平均注意方向再加上梯度懲罰。生成器的損失更簡單目標是讓生成樣本在critic那里的得分盡量高所以取負號。with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) fake_imgs tf.stop_gradient(fake_imgs) d_real critic(real_imgs, trainingTrue) d_fake critic(fake_imgs, trainingTrue) gp gradient_penalty(critic, real_imgs, fake_imgs) d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) lambda_gp * gp grads tape.gradient(d_loss, critic.trainable_variables) d_optimizer.apply_gradients(zip(grads, critic.trainable_variables))注意tf.stop_gradient(fake_imgs)這一行。如果不加GradientTape會把生成器的完整計算圖記錄下來。雖然我們只對critic參數(shù)求梯度但這部分記錄會白白消耗大量顯存訓練速度也會變慢。加上stop_gradient之后tape只追蹤從fake_imgs到d_fake這段計算內(nèi)存占用小很多。生成器的更新循環(huán)with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) d_fake critic(fake_imgs, trainingTrue) g_loss -tf.reduce_mean(d_fake) grads tape.gradient(g_loss, generator.trainable_variables) g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))整個訓練流程是先固定生成器更新critic n_critic次再固定critic更新生成器1次。critic的更新目標是拉開真實樣本得分和生成樣本得分的差距生成器的更新目標是拉高自己樣本的得分。兩者在d_fake這個點上形成對抗。3.3 訓練循環(huán)里的三個細節(jié)第一個細節(jié)是lambda的取值。WGAN-GP論文給出的經(jīng)典值就是10這個值在大多數(shù)任務(wù)上都不需要改。lambda太小時Lipschitz約束形同虛設(shè)訓練不穩(wěn)定lambda太大時critic會把幾乎所有注意力都放在滿足梯度約束上對真實樣本和生成樣本的區(qū)分能力下降生成圖像會變得模糊。如果調(diào)節(jié)后有類似問題優(yōu)先檢查它。第二個細節(jié)是critic的最后一層不要初始化偏置為0或權(quán)重過大。建議使用默認的Glorot均勻初始化如果發(fā)現(xiàn)訓練一開始critic輸出就很大可以考慮用小標準差初始化最后一層。第三個細節(jié)是d_loss的數(shù)值范圍。WGAN訓練中critic的loss是一個不斷變化的負數(shù)不會像傳統(tǒng)判別器那樣收斂到0附近。我見過很多人看到loss是負數(shù)就以為訓練出了問題其實這是正常的。WGAN-GP中d_loss穩(wěn)步下降、g_loss圍繞一個負值波動只能說明損失函數(shù)這一路沒有崩潰真正的判斷標準還是生成圖像本身。4. WGAN-Div從固定約束到可調(diào)散度4.1 WGAN-Div的原理與損失形式WGAN-Div站在WGAN-GP的肩膀上進一步把“梯度懲罰”從一種工程技巧上升成了理論上的“散度”定義。它引入兩個超參數(shù)k和p懲罰項寫成了k * mean(|grad_norm - k|^p)。這里要交代清楚一個細節(jié)WGAN-Div論文原文里的懲罰項寫法與工程代碼并不完全一致。原文更接近k * E[(k - grad_norm)^p]沒有先取絕對值當梯度范數(shù)大于k時懲罰項可能變號。實際測試中這種寫法在部分p取值下會引入不穩(wěn)定因素。所以我下面給出的代碼采用開源社區(qū)里更通用的穩(wěn)定版本對grad_norm - k取絕對值再求p次冪。如果你想嚴格對齊論文可以自己去掉絕對值對比效果。def wdiv_penalty(critic, real, fake, k1.0, p2.0): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3]) 1e-12) penalty k * tf.reduce_mean(tf.abs(norm - k) ** p) return penalty加1e-12是為了防止梯度范數(shù)為0時開根號出現(xiàn)數(shù)值不穩(wěn)。然后critic的損失函數(shù)變成d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) wdiv_penalty(critic, real_imgs, fake_imgs, k1.0, p2.0)乍一看和WGAN-GP幾乎一樣區(qū)別只在懲罰項的形式。WGAN-GP是強制梯度范數(shù)等于1WGAN-Div是強制梯度范數(shù)等于k并且用p次冪控制懲罰力度。當k1、p2時兩者非常接近可以互相替代。4.2 與WGAN-GP的代碼差異把WGAN-GP換成WGAN-Div代碼改動量其實非常小。生成器結(jié)構(gòu)、critic結(jié)構(gòu)、訓練循環(huán)框架全部可以復用只替換梯度懲罰函數(shù)。這也是我推薦在同一個工程里同時保留兩套實現(xiàn)的原因切換成本極低。WGAN-Div訓練循環(huán)中生成器的loss不變還是-mean(d_fake)。唯一要注意的是訓練critic時同樣要用stop_gradient隔離生成器擴大提高內(nèi)存效率。4.3 p值對訓練的影響從實際效果看p值比k值更值得玩味。p2時懲罰項對“梯度范數(shù)偏離k”施加的是二次懲罰類似WGAN-GPp增大到4或6意味著梯度范數(shù)偏差越大受到的懲罰越重。這意味著critic會更努力地讓梯度范數(shù)保持在k附近對局部梯度的約束更嚴格。論文里的實驗偏向于用較大的p比如p6認為這讓critic對空間各處的Lipschitz約束更均勻。我自己的測試結(jié)論是p2已經(jīng)夠用p6在某些數(shù)據(jù)集上能略微提升生成圖像的多樣性但代價是更容易出現(xiàn)數(shù)值波動尤其是訓練中期。p的調(diào)節(jié)不要一上來就拉滿先在p2跑通再逐檔對比。k值的作用更像一個“目標梯度范數(shù)”。k越小critic輸出對輸入的敏感度越低生成器收到的梯度相對溫和k越大critic可以更“尖銳”地響應(yīng)輸入變化生成器能獲得更強的修正信號但也更容易震蕩。我的經(jīng)驗是k取1或2都可以優(yōu)先從1開始。5. 實測對比與調(diào)參筆記5.1 兩個模型的訓練曲線差異把兩個模型都訓練20個epoch之后我記錄到一些有意思的現(xiàn)象。WGAN-GP的critic loss下降非常平穩(wěn)基本是一條斜率穩(wěn)定的曲線生成器loss則在一個負值附近小幅波動。WGAN-Div的曲線形態(tài)相似但絕對值大小和波幅會隨p變化p越大critic loss的數(shù)值抖動越明顯。這里必須強調(diào)不要試圖用critic loss的絕對值去橫向?qū)Ρ葍蓚€模型的優(yōu)劣。WGAN的loss是有意義的“分布距離估計值”但它受梯度懲罰項的影響較大改變罰項形式后數(shù)值就沒有可比性了。判斷模型好壞一看出圖效果二看是否出現(xiàn)模式崩塌三看訓練是否持續(xù)平穩(wěn)。5.2 生成質(zhì)量與穩(wěn)定性對比在MNIST上兩個模型最終生成的數(shù)字都清晰可辨幾乎看不出明顯差別。WGAN-GP從第8個epoch左右就能看到像樣的數(shù)字輪廓WGAN-Div在k1、p2時速度類似換成k1、p6后收斂略慢一點但生成樣本的多樣性感覺稍好一些。在CIFAR-10上兩個模型的差別開始顯現(xiàn)。WGAN-GP生成的圖像整體更“干凈”邊緣和色塊分明WGAN-Div在p較大時生成的圖像偶爾能保留更多細節(jié)紋理但也更容易出現(xiàn)局部顏色溢出。也就是說WGAN-Div的調(diào)節(jié)空間更大但它的最優(yōu)參數(shù)需要針對數(shù)據(jù)集單獨摸索不像WGAN-GP那樣“開箱即用”。這兩條路線我用表格簡單對比一下對比維度WGAN-GPWGAN-Div懲罰目標梯度范數(shù)等于1梯度范數(shù)等于k懲罰形式λ·E[(norm?1)2]k·E[關(guān)鍵超參λ默認10基本不用調(diào)k、p需要試驗理論統(tǒng)一度面向工程直觀從散度定義出發(fā)更一般訓練穩(wěn)定性很穩(wěn)適合首選穩(wěn)定但依賴參數(shù)選擇實現(xiàn)復雜度低極低改動一行適用場景絕大多數(shù)GAN任務(wù)想深挖理論或?qū)ふ叶鄻有允找?.3 關(guān)鍵超參調(diào)參心得學習率是這場實驗里影響最大的超參沒有之一。我試過把critic和生成器的學習率都調(diào)到3e-4訓練很快變得顛簸調(diào)回1e-4后馬上一路順風。WGAN系列對學習率比傳統(tǒng)GAN更敏感因為critic的梯度是要被“拿來算距離”的步子邁太大等于距離估錯。Batch size也不用太大。64是性價比很高的選擇。batch太大時梯度懲罰的插值樣本會更密集理論上更接近理論要求但實際中顯存壓力和訓練速度都會成問題收益卻不成比例。batch太小比如16critic會過度關(guān)注當前小批量中的個別樣本梯度噪聲偏大。n_critic在MNIST上試過1和3收斂速度差不太多在CIFAR-10上n_critic5比1更穩(wěn)。我的建議是復雜數(shù)據(jù)集用5簡單數(shù)據(jù)集可以縮到1或3省下來的算力拿去多跑幾個epoch更劃算。5.4 一個偷懶但有效的調(diào)試技巧訓練過程中一定要定期保存生成圖。我習慣每隔500個step把固定的一組噪聲向量丟進生成器生成一張對比圖保存下來。這樣跑完一遍訓練把這些圖按順序連起來看就是一個動畫訓練過程中有沒有崩潰、有沒有從噪聲慢慢變成圖像、有沒有突然模式崩塌一眼就能看清楚。只盯loss曲線是看不出這些信息的。我見過loss曲線一路下行但生成圖像全部是同一個數(shù)字的情況。模式崩塌時critic依然可以把真實樣本和“那一類”生成樣本區(qū)分開loss看起來一切正常。6. 問題排查與避坑記錄6.1 訓練崩壞的幾個典型原因癥狀可能的坑解決方向loss變成NaN學習率過高、梯度爆炸學習率降到1e-4以下檢查插值樣本是否有inf圖像長期模糊lambda太大 / 數(shù)據(jù)與輸出層不匹配lambda從10往下降或修正激活函數(shù)生成內(nèi)容單一模式崩塌減小batch、降低beta_1、增大latent_dimcritic loss快速歸零去掉sigmoid了嗎經(jīng)典GAN的判別頭混入了確認critic輸出層沒有sigmoid圖像顏色發(fā)灰數(shù)據(jù)范圍與激活函數(shù)不一致數(shù)據(jù)[-1,1]配tanh數(shù)據(jù)[0,1]配sigmoid6.2 網(wǎng)絡(luò)結(jié)構(gòu)相關(guān)的坑critic里加BN是我踩過最深的坑之一。最開始為了“讓訓練更穩(wěn)”按分類網(wǎng)絡(luò)的習慣在critic每個卷積后加了BN結(jié)果訓練時好時壞。后來查資料才發(fā)現(xiàn)梯度懲罰要求對單個輸入樣本求梯度BN卻讓輸出依賴于整個batch的統(tǒng)計量等于強行讓梯度“去個體化”。去掉BN后訓練立刻穩(wěn)定不少這也直接改變了我的網(wǎng)絡(luò)設(shè)計習慣。生成器里BN的training參數(shù)同樣是個容易出錯的點。如果你在訓練循環(huán)里調(diào)用generator(noise)時不帶training參數(shù)Keras默認按推理模式執(zhí)行BN不會更新。短期看loss還在變化但生成器內(nèi)部統(tǒng)計量一直沒動訓練效果會越來越差。正確做法是生成器在訓練階段顯式傳trainingTrue。6.3 數(shù)據(jù)與版本相關(guān)的坑圖像歸一化一定要在做batch之前完成。如果你用TensorFlow的Dataset API注意map里的歸一化操作不能在batch之后才做。我試過一次先batch再map歸一化結(jié)果歸一化作用的維度從(H, W, C)變成了(B, H, W, C)整個訓練損失數(shù)值都不對修了很久才定位到。版本兼容上TensorFlow 2.4之后tf.keras和獨立Keras的混用經(jīng)常出問題。我建議代碼里統(tǒng)一寫from tensorflow.keras import layers別一會import keras一會import tf.keras兩個庫的變量各自獨立加載權(quán)重和保存模型時會莫名報錯。我個人在實際操作中的體會是WGAN-GP和WGAN-Div不是競爭關(guān)系而是兩組可以隨時切換的“梯度約束方案”。如果手里有個圖像生成任務(wù)想快速出效果直接上WGAN-GP默認參數(shù)基本夠用如果愿意花時間調(diào)參想觀察多樣性變化再切到WGAN-Div去試不同的k和p。最后再分享一個小技巧實驗記錄時把每個step保存的生成圖統(tǒng)一命名成帶step編號的文件訓練完了用腳本拼成動畫你會非常直觀地看到模型是怎么從一片噪聲里慢慢長出結(jié)構(gòu)的——那種感覺比任何指標都讓人上癮。 WGAN-GP與WGAN-Div實戰(zhàn)：用Keras實現(xiàn)穩(wěn)定GAN訓練的梯度懲罰方案
如果你也經(jīng)歷過GAN訓練中那種玄學式的崩潰——生成器輸出一堆噪點、loss曲線像心電圖一樣亂跳、換一個隨機種子結(jié)果就天差地別——那么WGAN-GP和WGAN-Div這兩套方案應(yīng)該能幫你從泥潭里爬出來。我最近用Keras在MNIST和CIFAR-10上把這兩個模型完整搭了一遍都能穩(wěn)定生成圖像。這篇文章就把從原理到代碼再到調(diào)參的整個過程記錄下來給想動手復現(xiàn)的人一個可以直接抄作業(yè)的參考。先說明一個前提下面所有代碼基于TensorFlow 2.x內(nèi)置的Keras。如果你還在用獨立的Keras 2.xAPI基本通用但GradientTape那一套只存在于tf.keras里建議直接切換到TensorFlow 2.x省掉版本兼容的折騰。1. 為什么繞不開WGAN這條路線1.1 經(jīng)典GAN的病灶JS散度在偷懶傳統(tǒng)GAN里判別器輸出的是一個概率用交叉熵做loss本質(zhì)上是在最小化生成分布和真實分布之間的JS散度。這個度量方式有個致命問題當兩個分布的重疊區(qū)域非常小的時候JS散度會趨于一個常數(shù)log2梯度幾乎為0。舉個例子。假設(shè)真實數(shù)據(jù)分布是二維平面上的一條細線生成器剛開始只能撒出一些離散的噪點。這兩組點在絕大多數(shù)位置都不重合判別器很容易就把它們分得明明白白。此時生成器幾乎收不到任何有用的梯度信號不管怎么調(diào)參數(shù)輸出的還是那堆噪點。這就是大家常說的“判別器太強生成器被錘死”。實際操作中的表現(xiàn)就是判別器的loss快速降到接近0生成器的loss卻不怎么動生成的圖片越來越像噪聲。很多人第一反應(yīng)是調(diào)判別器的學習率但根子在于度量方式本身就給不出有效的“距離感”。1.2 Wasserstein距離換了個度量方式WGAN的核心改動是把判別器換成了“critic”輸出不再是一個概率而是一個實數(shù)得分。它優(yōu)化的目標也從JS散度換成Wasserstein距離也就是常說的“推土機距離”。這個距離的直覺很好理解假設(shè)真實分布是一堆土生成分布是另一堆土Wasserstein距離就是把這堆土搬到另一堆土的最小運輸成本。即使兩堆土完全不重疊只要它們位置不同這個成本就是一個正的有限值梯度就不會消失。這就解決了經(jīng)典GAN里“距離太遠就無梯度”的問題。所以WGAN的critic要做的事情不是判斷“真”“假”而是給真實樣本打高分給生成樣本打低分。對應(yīng)到代碼里就是去掉最后一層的sigmoid輸出一個raw score。這是WGAN與傳統(tǒng)GAN最直觀的區(qū)別。1.3 權(quán)重裁剪的副作用WGAN為了保證Wasserstein距離能算得下去要求critic的函數(shù)滿足1-Lipschitz條件也就是輸出對輸入的改變不能過于敏感。論文作者最初用的方法是權(quán)重裁剪把critic的每一層權(quán)重強行限制在一個小范圍內(nèi)比如[-0.01, 0.01]。這個方法實現(xiàn)起來非常簡單但副作用很明顯。權(quán)重被強行裁剪后critic的表達能力被壓縮得很厲害很多情況下它只能學到非常粗糙的判別函數(shù)。我一開始照著原始WGAN代碼跑發(fā)現(xiàn)生成的圖像總有一種“模糊感”細節(jié)上不去。后來才知道這是權(quán)重裁剪帶來的典型問題critic太弱給生成器提供的梯度質(zhì)量也差。WGAN-GP和WGAN-Div就是在這個基礎(chǔ)上提出的兩種改進。它們的核心思路一模一樣不再裁剪權(quán)重而是通過給損失函數(shù)加一個“梯度懲罰項”直接約束critic對輸入的梯度從而滿足Lipschitz條件。2. 實驗整體設(shè)計先定結(jié)構(gòu)再談?chuàng)p失2.1 數(shù)據(jù)準備與歸一化我用兩個數(shù)據(jù)集做驗證MNIST用來快速驗證收斂邏輯CIFAR-10用來驗證彩色、復雜圖像上的生成效果。兩個數(shù)據(jù)集都縮放到[-1, 1]區(qū)間生成器最后一層用tanh激活輸出自然落在[-1, 1]內(nèi)正好對應(yīng)輸入的數(shù)據(jù)范圍。這里有個容易忽略的細節(jié)如果你把圖像歸一化到[0, 1]生成器最后一層還在用tanh輸出那生成圖像的像素會被強制壓到[-1, 1]看起來會“發(fā)灰”。反過來如果數(shù)據(jù)在[-1, 1]內(nèi)生成器卻用sigmoid輸出那圖像永遠達不到最大亮度和對比度。所以訓練前先確認好數(shù)據(jù)的范圍與輸出層的激活函數(shù)。2.2 生成器結(jié)構(gòu)生成器沿用DCGAN那套結(jié)構(gòu)先通過全連接層把隨機噪聲映射成一個小尺寸的特征圖再用轉(zhuǎn)置卷積逐步上采樣到目標尺寸。我在CIFAR-10上用32×32分辨率從8×8的特征圖開始連續(xù)兩次stride為2的轉(zhuǎn)置卷積到16×16、32×32最后一層改用普通卷積輸出3個通道。def build_generator(latent_dim128): noise Input(shape(latent_dim,)) x Dense(8 * 8 * 256, use_biasFalse)(noise) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Reshape((8, 8, 256))(x) x Conv2DTranspose(128, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2DTranspose(64, (4, 4), strides(2, 2), paddingsame, use_biasFalse)(x) x BatchNormalization()(x) x LeakyReLU(alpha0.2)(x) x Conv2D(3, (3, 3), paddingsame, activationtanh)(x) return Model(noise, x)生成器里每一層都用LeakyReLUalpha取0.2。這個選擇是為了防止ReLU死掉——GAN訓練過程中生成器的輸入分布變化很大ReLU一旦被壓到負區(qū)間就可能一直輸出0導致梯度斷流。LeakyReLU可以保證梯度至少有一條通道能流回去。BN層在生成器里很重要能顯著穩(wěn)定訓練。但注意調(diào)用生成器時如果處于訓練階段一定要寫generator(noise, trainingTrue)讓BN更新批次統(tǒng)計量。生成圖像做推理時再用trainingFalse否則你會發(fā)現(xiàn)同一個噪聲向量每次生成的結(jié)果不一樣原因是BN在訓練模式下用了當前batch的均值方差。2.3 Critic結(jié)構(gòu)critic的結(jié)構(gòu)與生成器鏡像但有幾處關(guān)鍵差異不使用BN不使用sigmoid最后一層直接輸出一個標量。def build_critic(img_shape(32, 32, 3)): img Input(shapeimg_shape) x Conv2D(64, (4, 4), strides(2, 2), paddingsame)(img) x LeakyReLU(alpha0.2)(x) x Conv2D(128, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Conv2D(256, (4, 4), strides(2, 2), paddingsame)(x) x LeakyReLU(alpha0.2)(x) x Flatten()(x) x Dropout(0.2)(x) out Dense(1)(x) return Model(img, out)為什么critic不用BN梯度懲罰的計算要求在插值樣本上對輸入求導而BN在訓練時引入批內(nèi)統(tǒng)計量使輸出不再只依賴當前單個輸入相當于在批維度上產(chǎn)生了信息泄露。這樣一來“對單個輸入的梯度”就變得不再純粹懲罰項也會受到污染。很多第一次復現(xiàn)WGAN-GP的人把BN留在critic里結(jié)果訓練不穩(wěn)定多半是這個原因。如果卷積下采樣想保留更多信息也可以不加Dropout。我加了一層Dropout(0.2)是防止critic對訓練集過度敏感實測對穩(wěn)定性有一點幫助但不是必須的。2.4 訓練超參整套實驗的關(guān)鍵超參如下latent_dim 128batch_size 64critic更新次數(shù)n_critic 5MNIST上我試過1也能收斂優(yōu)化器Adamlr 1e-4beta_1 0.5WGAN-GP的lambda 10WGAN-Div的k 1p 2起步Adam的beta_1取0.5而不是默認的0.9這里說下原因。GAN的loss曲面本身波動劇烈如果動量累積過大梯度更新方向會“剎不住車”導致訓練震蕩。beta_1降低以后動量衰減更快更新方向更追隨當前梯度穩(wěn)定性會好很多。n_critic是一個可以調(diào)的空閑參數(shù)。理論上critic需要先訓練充分才能給生成器提供準確的Wasserstein梯度。但如果critic太強、生成器太弱同樣會出現(xiàn)生成器完全跟不上的問題。我在小數(shù)據(jù)集上的經(jīng)驗是n_critic1到5都能跑數(shù)據(jù)集越小、越簡單n_critic就可以越小。3. WGAN-GP梯度懲罰的完整實現(xiàn)3.1 插值樣本與梯度懲罰的計算WGAN-GP的核心是構(gòu)造一個“插值樣本”讓critic在這個樣本上的梯度范數(shù)盡量接近1。所謂插值樣本就是從一對真實樣本和生成樣本的連線上隨機取一個點alpha * real (1 - alpha) * fakealpha是[0,1]間的隨機數(shù)。為什么要在連線上采樣因為Lipschitz條件要求函數(shù)的梯度在全空間都有界但真正計算所有位置不可行。經(jīng)驗上真實樣本和生成樣本中間過渡區(qū)域是最容易違反Lipschitz條件的地方在連線上做懲罰性價比最高。這個思想不是嚴格的數(shù)學證明但實驗效果非常好。def gradient_penalty(critic, real, fake): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3])) gp tf.reduce_mean((norm - 1.0) ** 2) return gp這里有個初學者容易出錯的地方tf.GradientTape()里除了要計算critic的輸出還要顯式調(diào)用tape.watch(interp)否則TensorFlow不會自動記錄對interp的梯度。很多報錯“gradient is None”都是因為這個。3.2 損失函數(shù)與訓練循環(huán)WGAN-GP中critic的損失不是交叉熵而是三部分之和對真實樣本的輸出取平均對生成樣本的輸出取平均注意方向再加上梯度懲罰。生成器的損失更簡單目標是讓生成樣本在critic那里的得分盡量高所以取負號。with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) fake_imgs tf.stop_gradient(fake_imgs) d_real critic(real_imgs, trainingTrue) d_fake critic(fake_imgs, trainingTrue) gp gradient_penalty(critic, real_imgs, fake_imgs) d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) lambda_gp * gp grads tape.gradient(d_loss, critic.trainable_variables) d_optimizer.apply_gradients(zip(grads, critic.trainable_variables))注意tf.stop_gradient(fake_imgs)這一行。如果不加GradientTape會把生成器的完整計算圖記錄下來。雖然我們只對critic參數(shù)求梯度但這部分記錄會白白消耗大量顯存訓練速度也會變慢。加上stop_gradient之后tape只追蹤從fake_imgs到d_fake這段計算內(nèi)存占用小很多。生成器的更新循環(huán)with tf.GradientTape() as tape: fake_imgs generator(noise, trainingTrue) d_fake critic(fake_imgs, trainingTrue) g_loss -tf.reduce_mean(d_fake) grads tape.gradient(g_loss, generator.trainable_variables) g_optimizer.apply_gradients(zip(grads, generator.trainable_variables))整個訓練流程是先固定生成器更新critic n_critic次再固定critic更新生成器1次。critic的更新目標是拉開真實樣本得分和生成樣本得分的差距生成器的更新目標是拉高自己樣本的得分。兩者在d_fake這個點上形成對抗。3.3 訓練循環(huán)里的三個細節(jié)第一個細節(jié)是lambda的取值。WGAN-GP論文給出的經(jīng)典值就是10這個值在大多數(shù)任務(wù)上都不需要改。lambda太小時Lipschitz約束形同虛設(shè)訓練不穩(wěn)定lambda太大時critic會把幾乎所有注意力都放在滿足梯度約束上對真實樣本和生成樣本的區(qū)分能力下降生成圖像會變得模糊。如果調(diào)節(jié)后有類似問題優(yōu)先檢查它。第二個細節(jié)是critic的最后一層不要初始化偏置為0或權(quán)重過大。建議使用默認的Glorot均勻初始化如果發(fā)現(xiàn)訓練一開始critic輸出就很大可以考慮用小標準差初始化最后一層。第三個細節(jié)是d_loss的數(shù)值范圍。WGAN訓練中critic的loss是一個不斷變化的負數(shù)不會像傳統(tǒng)判別器那樣收斂到0附近。我見過很多人看到loss是負數(shù)就以為訓練出了問題其實這是正常的。WGAN-GP中d_loss穩(wěn)步下降、g_loss圍繞一個負值波動只能說明損失函數(shù)這一路沒有崩潰真正的判斷標準還是生成圖像本身。4. WGAN-Div從固定約束到可調(diào)散度4.1 WGAN-Div的原理與損失形式WGAN-Div站在WGAN-GP的肩膀上進一步把“梯度懲罰”從一種工程技巧上升成了理論上的“散度”定義。它引入兩個超參數(shù)k和p懲罰項寫成了k * mean(|grad_norm - k|^p)。這里要交代清楚一個細節(jié)WGAN-Div論文原文里的懲罰項寫法與工程代碼并不完全一致。原文更接近k * E[(k - grad_norm)^p]沒有先取絕對值當梯度范數(shù)大于k時懲罰項可能變號。實際測試中這種寫法在部分p取值下會引入不穩(wěn)定因素。所以我下面給出的代碼采用開源社區(qū)里更通用的穩(wěn)定版本對grad_norm - k取絕對值再求p次冪。如果你想嚴格對齊論文可以自己去掉絕對值對比效果。def wdiv_penalty(critic, real, fake, k1.0, p2.0): alpha tf.random.uniform((tf.shape(real)[0], 1, 1, 1), 0.0, 1.0) interp alpha * real (1.0 - alpha) * fake with tf.GradientTape() as tape: tape.watch(interp) pred critic(interp, trainingTrue) grad tape.gradient(pred, interp) norm tf.sqrt(tf.reduce_sum(tf.square(grad), axis[1, 2, 3]) 1e-12) penalty k * tf.reduce_mean(tf.abs(norm - k) ** p) return penalty加1e-12是為了防止梯度范數(shù)為0時開根號出現(xiàn)數(shù)值不穩(wěn)。然后critic的損失函數(shù)變成d_loss tf.reduce_mean(d_fake) - tf.reduce_mean(d_real) wdiv_penalty(critic, real_imgs, fake_imgs, k1.0, p2.0)乍一看和WGAN-GP幾乎一樣區(qū)別只在懲罰項的形式。WGAN-GP是強制梯度范數(shù)等于1WGAN-Div是強制梯度范數(shù)等于k并且用p次冪控制懲罰力度。當k1、p2時兩者非常接近可以互相替代。4.2 與WGAN-GP的代碼差異把WGAN-GP換成WGAN-Div代碼改動量其實非常小。生成器結(jié)構(gòu)、critic結(jié)構(gòu)、訓練循環(huán)框架全部可以復用只替換梯度懲罰函數(shù)。這也是我推薦在同一個工程里同時保留兩套實現(xiàn)的原因切換成本極低。WGAN-Div訓練循環(huán)中生成器的loss不變還是-mean(d_fake)。唯一要注意的是訓練critic時同樣要用stop_gradient隔離生成器擴大提高內(nèi)存效率。4.3 p值對訓練的影響從實際效果看p值比k值更值得玩味。p2時懲罰項對“梯度范數(shù)偏離k”施加的是二次懲罰類似WGAN-GPp增大到4或6意味著梯度范數(shù)偏差越大受到的懲罰越重。這意味著critic會更努力地讓梯度范數(shù)保持在k附近對局部梯度的約束更嚴格。論文里的實驗偏向于用較大的p比如p6認為這讓critic對空間各處的Lipschitz約束更均勻。我自己的測試結(jié)論是p2已經(jīng)夠用p6在某些數(shù)據(jù)集上能略微提升生成圖像的多樣性但代價是更容易出現(xiàn)數(shù)值波動尤其是訓練中期。p的調(diào)節(jié)不要一上來就拉滿先在p2跑通再逐檔對比。k值的作用更像一個“目標梯度范數(shù)”。k越小critic輸出對輸入的敏感度越低生成器收到的梯度相對溫和k越大critic可以更“尖銳”地響應(yīng)輸入變化生成器能獲得更強的修正信號但也更容易震蕩。我的經(jīng)驗是k取1或2都可以優(yōu)先從1開始。5. 實測對比與調(diào)參筆記5.1 兩個模型的訓練曲線差異把兩個模型都訓練20個epoch之后我記錄到一些有意思的現(xiàn)象。WGAN-GP的critic loss下降非常平穩(wěn)基本是一條斜率穩(wěn)定的曲線生成器loss則在一個負值附近小幅波動。WGAN-Div的曲線形態(tài)相似但絕對值大小和波幅會隨p變化p越大critic loss的數(shù)值抖動越明顯。這里必須強調(diào)不要試圖用critic loss的絕對值去橫向?qū)Ρ葍蓚€模型的優(yōu)劣。WGAN的loss是有意義的“分布距離估計值”但它受梯度懲罰項的影響較大改變罰項形式后數(shù)值就沒有可比性了。判斷模型好壞一看出圖效果二看是否出現(xiàn)模式崩塌三看訓練是否持續(xù)平穩(wěn)。5.2 生成質(zhì)量與穩(wěn)定性對比在MNIST上兩個模型最終生成的數(shù)字都清晰可辨幾乎看不出明顯差別。WGAN-GP從第8個epoch左右就能看到像樣的數(shù)字輪廓WGAN-Div在k1、p2時速度類似換成k1、p6后收斂略慢一點但生成樣本的多樣性感覺稍好一些。在CIFAR-10上兩個模型的差別開始顯現(xiàn)。WGAN-GP生成的圖像整體更“干凈”邊緣和色塊分明WGAN-Div在p較大時生成的圖像偶爾能保留更多細節(jié)紋理但也更容易出現(xiàn)局部顏色溢出。也就是說WGAN-Div的調(diào)節(jié)空間更大但它的最優(yōu)參數(shù)需要針對數(shù)據(jù)集單獨摸索不像WGAN-GP那樣“開箱即用”。這兩條路線我用表格簡單對比一下對比維度WGAN-GPWGAN-Div懲罰目標梯度范數(shù)等于1梯度范數(shù)等于k懲罰形式λ·E[(norm?1)2]k·E[關(guān)鍵超參λ默認10基本不用調(diào)k、p需要試驗理論統(tǒng)一度面向工程直觀從散度定義出發(fā)更一般訓練穩(wěn)定性很穩(wěn)適合首選穩(wěn)定但依賴參數(shù)選擇實現(xiàn)復雜度低極低改動一行適用場景絕大多數(shù)GAN任務(wù)想深挖理論或?qū)ふ叶鄻有允找?.3 關(guān)鍵超參調(diào)參心得學習率是這場實驗里影響最大的超參沒有之一。我試過把critic和生成器的學習率都調(diào)到3e-4訓練很快變得顛簸調(diào)回1e-4后馬上一路順風。WGAN系列對學習率比傳統(tǒng)GAN更敏感因為critic的梯度是要被“拿來算距離”的步子邁太大等于距離估錯。Batch size也不用太大。64是性價比很高的選擇。batch太大時梯度懲罰的插值樣本會更密集理論上更接近理論要求但實際中顯存壓力和訓練速度都會成問題收益卻不成比例。batch太小比如16critic會過度關(guān)注當前小批量中的個別樣本梯度噪聲偏大。n_critic在MNIST上試過1和3收斂速度差不太多在CIFAR-10上n_critic5比1更穩(wěn)。我的建議是復雜數(shù)據(jù)集用5簡單數(shù)據(jù)集可以縮到1或3省下來的算力拿去多跑幾個epoch更劃算。5.4 一個偷懶但有效的調(diào)試技巧訓練過程中一定要定期保存生成圖。我習慣每隔500個step把固定的一組噪聲向量丟進生成器生成一張對比圖保存下來。這樣跑完一遍訓練把這些圖按順序連起來看就是一個動畫訓練過程中有沒有崩潰、有沒有從噪聲慢慢變成圖像、有沒有突然模式崩塌一眼就能看清楚。只盯loss曲線是看不出這些信息的。我見過loss曲線一路下行但生成圖像全部是同一個數(shù)字的情況。模式崩塌時critic依然可以把真實樣本和“那一類”生成樣本區(qū)分開loss看起來一切正常。6. 問題排查與避坑記錄6.1 訓練崩壞的幾個典型原因癥狀可能的坑解決方向loss變成NaN學習率過高、梯度爆炸學習率降到1e-4以下檢查插值樣本是否有inf圖像長期模糊lambda太大 / 數(shù)據(jù)與輸出層不匹配lambda從10往下降或修正激活函數(shù)生成內(nèi)容單一模式崩塌減小batch、降低beta_1、增大latent_dimcritic loss快速歸零去掉sigmoid了嗎經(jīng)典GAN的判別頭混入了確認critic輸出層沒有sigmoid圖像顏色發(fā)灰數(shù)據(jù)范圍與激活函數(shù)不一致數(shù)據(jù)[-1,1]配tanh數(shù)據(jù)[0,1]配sigmoid6.2 網(wǎng)絡(luò)結(jié)構(gòu)相關(guān)的坑critic里加BN是我踩過最深的坑之一。最開始為了“讓訓練更穩(wěn)”按分類網(wǎng)絡(luò)的習慣在critic每個卷積后加了BN結(jié)果訓練時好時壞。后來查資料才發(fā)現(xiàn)梯度懲罰要求對單個輸入樣本求梯度BN卻讓輸出依賴于整個batch的統(tǒng)計量等于強行讓梯度“去個體化”。去掉BN后訓練立刻穩(wěn)定不少這也直接改變了我的網(wǎng)絡(luò)設(shè)計習慣。生成器里BN的training參數(shù)同樣是個容易出錯的點。如果你在訓練循環(huán)里調(diào)用generator(noise)時不帶training參數(shù)Keras默認按推理模式執(zhí)行BN不會更新。短期看loss還在變化但生成器內(nèi)部統(tǒng)計量一直沒動訓練效果會越來越差。正確做法是生成器在訓練階段顯式傳trainingTrue。6.3 數(shù)據(jù)與版本相關(guān)的坑圖像歸一化一定要在做batch之前完成。如果你用TensorFlow的Dataset API注意map里的歸一化操作不能在batch之后才做。我試過一次先batch再map歸一化結(jié)果歸一化作用的維度從(H, W, C)變成了(B, H, W, C)整個訓練損失數(shù)值都不對修了很久才定位到。版本兼容上TensorFlow 2.4之后tf.keras和獨立Keras的混用經(jīng)常出問題。我建議代碼里統(tǒng)一寫from tensorflow.keras import layers別一會import keras一會import tf.keras兩個庫的變量各自獨立加載權(quán)重和保存模型時會莫名報錯。我個人在實際操作中的體會是WGAN-GP和WGAN-Div不是競爭關(guān)系而是兩組可以隨時切換的“梯度約束方案”。如果手里有個圖像生成任務(wù)想快速出效果直接上WGAN-GP默認參數(shù)基本夠用如果愿意花時間調(diào)參想觀察多樣性變化再切到WGAN-Div去試不同的k和p。最后再分享一個小技巧實驗記錄時把每個step保存的生成圖統(tǒng)一命名成帶step編號的文件訓練完了用腳本拼成動畫你會非常直觀地看到模型是怎么從一片噪聲里慢慢長出結(jié)構(gòu)的——那種感覺比任何指標都讓人上癮。