簡(jiǎn)介壓縮包內(nèi)含一套去噪擴(kuò)散概率模型Diffusion Model的Python實(shí)現(xiàn)適合深度學(xué)習(xí)、計(jì)算機(jī)視覺(jué)方向的學(xué)生與算法工程師用于圖像生成實(shí)驗(yàn)或二次開(kāi)發(fā)。代碼覆蓋模型核心組件、訓(xùn)練工具與數(shù)據(jù)集加載邏輯并針對(duì)CelebA-HQ、LSUN、CIFAR等常用數(shù)據(jù)集提供運(yùn)行腳本可幫助理解擴(kuò)散模型的訓(xùn)練與采樣流程。整個(gè)zip共20個(gè)文件以17個(gè).py代碼文件為主同時(shí)包含README說(shuō)明、requirements依賴清單和png效果圖總大小約985KB量級(jí)輕量、結(jié)構(gòu)直觀。目前已有325人瀏覽學(xué)習(xí)適合作為課程設(shè)計(jì)、期末大作業(yè)或畢業(yè)設(shè)計(jì)的參考基線代碼采用參數(shù)化編程注釋明細(xì)便于修改超參數(shù)和替換數(shù)據(jù)集結(jié)合附帶的預(yù)覽圖與說(shuō)明文檔可快速驗(yàn)證生成效果。1. diffusion去噪擴(kuò)散概率模型這份Python代碼包里藏著一個(gè)能訓(xùn)練的DDPM很多人第一次接觸diffusion去噪擴(kuò)散概率模型是因?yàn)镾table Diffusion的生圖效果但真要復(fù)現(xiàn)一個(gè)能訓(xùn)練、能采樣、能出圖的擴(kuò)散模型大多數(shù)時(shí)候卡在論文公式和工程實(shí)現(xiàn)的縫隙里。這份資源不是零散的筆記而是一個(gè)按可運(yùn)行方式組織的Python工程入口腳本、擴(kuò)散工具、U-Net模型和數(shù)據(jù)接口都齊了覆蓋CIFAR、CelebA-HQ、LSUN三個(gè)數(shù)據(jù)集。適合做課程設(shè)計(jì)、期末大作業(yè)和畢業(yè)設(shè)計(jì)的計(jì)算機(jī)、電子信息、數(shù)學(xué)方向?qū)W生也適合想從零跑通一次DDPM的從業(yè)者直接照著入口腳本改參數(shù)先跑通再回頭讀公式比從零寫整個(gè)框架省力得多。就算你之前習(xí)慣在Matlab里做算法仿真這種參數(shù)化組織和清晰注釋也能讓你快速把思路遷到自己的工具鏈上。2. 從加噪到去噪DDPM前向過(guò)程與反向訓(xùn)練的核心公式2.1 前向過(guò)程為什么說(shuō)擴(kuò)散就是把圖像一步步抹成噪聲拿到這個(gè)包第一件事別急著跑訓(xùn)練先把diffusion_utils.py里的變量名和公式對(duì)上。DDPM的前向過(guò)程是一條馬爾可夫鏈給定一張真實(shí)圖像 x_0每一步按 x_t sqrt(1 - beta_t) * x_{t-1} sqrt(beta_t) * z_t 注入高斯噪聲 z_t。這里 beta_t 是預(yù)設(shè)的噪聲權(quán)重原論文用線性調(diào)度從 beta_1 1e-4 一路加到 beta_T 0.02時(shí)間步 T 默認(rèn)取 1000。也就是說(shuō)模型看到的不是一步步變模糊而是一張圖在 1000 步內(nèi)被逐漸淹沒(méi)成標(biāo)準(zhǔn)正態(tài)噪聲的過(guò)程。你可能會(huì)問(wèn)為什么要搞 1000 步這么多因?yàn)橹挥邪鸭釉脒^(guò)程切得足夠細(xì)反向去噪時(shí)每一步的改動(dòng)才足夠小神經(jīng)網(wǎng)絡(luò)才容易擬合。這個(gè)直覺(jué)貫穿整個(gè)DDPM設(shè)計(jì)。實(shí)際代碼里不會(huì)真的逐步迭代 1000 次來(lái)制造訓(xùn)練樣本而是利用高斯噪聲的可疊加性直接寫出從 x_0 到任意時(shí)刻 x_t 的閉合表達(dá)式# q_sample: 直接從 x_0 跳到任意時(shí)刻 t 的帶噪樣本避免逐步迭代 def q_sample(x_start, t, noise, alphas_cumprod): # alphas_cumprod: 由 (1 - beta_t) 連乘得到的累計(jì)序列長(zhǎng)度等于 T sqrt_acp torch.sqrt(alphas_cumprod[t])[:, None, None, None] sqrt_one_minus_acp torch.sqrt(1.0 - alphas_cumprod[t])[:, None, None, None] return sqrt_acp * x_start sqrt_one_minus_acp * noise這個(gè)函數(shù)是 diffusion_utils.py 里最底層的一段邏輯我用 PyTorch 寫法示意資源包里的 TensorFlow 實(shí)現(xiàn)變量名和公式一一對(duì)應(yīng)。q_sample 存在的意義是訓(xùn)練時(shí)不用從 t0 迭代到 t999而是隨機(jī)抽一個(gè)時(shí)刻 t一步算出 x_t。否則每個(gè) batch 都跑 1000 步前向訓(xùn)練速度會(huì)慢到完全沒(méi)法用。參數(shù)上要注意三點(diǎn)第一alphas_cumprod 是提前算好的長(zhǎng)度為 T 的向量別在訓(xùn)練循環(huán)里動(dòng)態(tài)計(jì)算第二t 通常按 batch 里的每個(gè)樣本獨(dú)立隨機(jī)采樣形狀要和 batch 維對(duì)齊第三noise 要每個(gè)樣本單獨(dú)采樣不能把同一個(gè)噪聲廣播到整個(gè) batch否則模型會(huì)學(xué)到樣本間的偽相關(guān)性。2.2 訓(xùn)練目標(biāo)預(yù)測(cè)噪聲比預(yù)測(cè)圖像更好收斂如果把訓(xùn)練目標(biāo)定成直接預(yù)測(cè) x_0早期時(shí)間步的 x_t 里幾乎全是噪聲回歸目標(biāo)方差巨大網(wǎng)絡(luò)很容易陷在模糊平均里。DDPM 的關(guān)鍵改動(dòng)是把目標(biāo)換成預(yù)測(cè)這個(gè)時(shí)刻被注入的噪聲。網(wǎng)絡(luò)不再猜完整圖像只回答我手上這團(tuán)帶噪圖里混了什么噪聲這件事在每一個(gè)時(shí)間步上都是良定義的。訓(xùn)練損失也因此異常簡(jiǎn)單就是一個(gè) MSE# p_losses: 訓(xùn)練時(shí)只做一件事讓預(yù)測(cè)噪聲逼近真實(shí)噪聲 def p_losses(denoise_fn, x_start, t, noise): x_noisy q_sample(x_startx_start, tt, noisenoise) # 先加噪 noise_pred denoise_fn(x_noisy, t) # U-Net 輸出預(yù)測(cè)噪聲 loss F.mse_loss(noise, noise_pred) # 與真實(shí)噪聲做 MSE return lossdenoise_fn 對(duì)應(yīng) models/nn.py 里的 U-Net輸入是帶噪圖像和時(shí)間步 tt 先經(jīng)過(guò)正弦位置編碼再和圖像特征在通道維度上融合。為什么用 MSE 不用 L1DDPM 原文在噪聲預(yù)測(cè)上用的就是 MSE配合 L2 損失梯度在整個(gè)時(shí)間步上更平滑訓(xùn)練過(guò)程更平穩(wěn)。訓(xùn)練時(shí)每個(gè)樣本隨機(jī)抽一個(gè) t等于讓同一個(gè)網(wǎng)絡(luò)同時(shí)學(xué)會(huì)去一點(diǎn)噪和去很多噪兩件事推理時(shí)再把它們串成一個(gè)從純?cè)肼暤綀D像的循環(huán)。實(shí)際訓(xùn)練中 loss 收斂到接近噪聲方差附近就可以考慮采樣本看看不必死等它降到 0——它本來(lái)就降不到 0。2.3 反向采樣從純?cè)肼曋鸩竭€原圖像的迭代循環(huán)推理階段和訓(xùn)練完全反過(guò)來(lái)從 x_T ~ N(0, I) 開(kāi)始按照逆過(guò)程 p(x_{t-1} | x_t) 一級(jí)一級(jí)去噪。DDPM 采樣公式里給定網(wǎng)絡(luò)預(yù)測(cè)的噪聲 epsx_{t-1} 的均值可以解析地算出來(lái)方差則是預(yù)先算好的后驗(yàn)方差。采樣核心代碼通常是這個(gè)樣子# p_sample: 單步反向去噪從 x_t 還原出 x_{t-1} def p_sample(model, x_t, t, alphas, alphas_cumprod, alphas_cumprod_prev): eps model(x_t, t) mu (x_t - (1 - alphas[t]) / torch.sqrt(1 - alphas_cumprod[t]) * eps) / torch.sqrt(alphas[t]) sigma torch.sqrt((1 - alphas_cumprod_prev[t]) / (1 - alphas_cumprod[t]) * (1 - alphas[t])) return mu sigma * torch.randn_like(x_t) # 最后一步 t0 時(shí) sigma 應(yīng)強(qiáng)制為 0mu 公式直接來(lái)自 DDPM 論文的采樣過(guò)程分子里的 (1 - alphas[t]) 是當(dāng)前步的噪聲權(quán)重分母里的 sqrt(1 - alphas_cumprod[t]) 用來(lái)把累積噪聲歸一化sigma 是后驗(yàn)方差在最后一個(gè)時(shí)間步 t0 時(shí)沒(méi)有隨機(jī)噪聲必須把 sigma 強(qiáng)制置 0否則生成結(jié)果會(huì)帶一層甩不掉的顆粒感。這段邏輯在源碼包里一般寫在 diffusion_utils.py 的 p_mean_variance 附近變量名可能是 alphas_cumprod_prev 和 posterior_variance思路完全一致。整個(gè)采樣循環(huán)就是從 T 遞減到 0把上一輪輸出當(dāng)成下一輪輸入最后拿到的 x_0 就是生成圖。2.4 調(diào)度參數(shù)與時(shí)間索引兩個(gè)經(jīng)常被忽略的細(xì)節(jié)很多人在小圖上跑通了一次換數(shù)據(jù)集就翻車問(wèn)題常出在 beta 調(diào)度和時(shí)間索引上。beta 調(diào)度不只有 linear 一種后續(xù)工作里 cosine 調(diào)度也很常見(jiàn)它在中間段加噪更緩和高分辨率數(shù)據(jù)上更容易收斂。如果你在腳本參數(shù)里看到 beta_schedule 選項(xiàng)默認(rèn) linear 可以先不動(dòng)但要知道把 T 從 1000 改成 500 時(shí)alphas_cumprod 長(zhǎng)度必須同步變否則索引越界或采樣公式錯(cuò)位。另一個(gè)細(xì)節(jié)是 alphas_cumprod_prev 的移位處理。采樣時(shí)計(jì)算 t-1 步的累積量一般用 prev cat([1, alphas_cumprod[:-1]]) 這種寫法讓 t0 時(shí)有 prev1這樣最后一個(gè)時(shí)間步才能正確終止。很多魔改代碼丟失這個(gè) shift生成結(jié)果會(huì)整體偏色。遇到這類問(wèn)題先打印 alphas_cumprod 最后幾個(gè)值和 prev 的前幾個(gè)值一眼就能看出錯(cuò)沒(méi)錯(cuò)。2.5 為什么這份代碼用TensorFlow實(shí)現(xiàn)對(duì)復(fù)現(xiàn)的影響資源包里的 diffusion_tf、tpu_utils.py、run_*.py 都指向 TensorFlow 實(shí)現(xiàn)風(fēng)格這也是 DDPM 早期開(kāi)源生態(tài)的一個(gè)特點(diǎn)不少論文官方代碼就是 TF后來(lái)大家才陸續(xù)移植到 PyTorch。對(duì)復(fù)現(xiàn)來(lái)說(shuō)這不是阻礙擴(kuò)散模型的數(shù)學(xué)和網(wǎng)絡(luò)架構(gòu)都是通用的TF 版本反而有一個(gè)好處就是能直接沿用當(dāng)年的 TPU 訓(xùn)練配置。tpu_utils.py 里那些分布式參數(shù)在單卡 GPU 上多數(shù)用不上但它的存在提醒你這套代碼是按大規(guī)模訓(xùn)練設(shè)計(jì)的不是玩具。如果你本地只有一張消費(fèi)級(jí) GPU建議把它當(dāng)參考實(shí)現(xiàn)重點(diǎn)讀 diffusion_utils.py 和 models/nn.py 的邏輯。若實(shí)在不習(xí)慣 TF把這兩處改寫成 PyTorch 是常規(guī)遷移路徑先跑通小數(shù)據(jù)集小分辨率確認(rèn) loss 和采樣行為正常再往上擴(kuò)規(guī)模。代碼里 TensorFlow 的 Eager 模式和 PyTorch 的 eager 語(yǔ)義高度接近逐個(gè)函數(shù)對(duì)照搬就行。3. 源碼包結(jié)構(gòu)拆解從文件組織看懂一條完整訓(xùn)練鏈路3.1 解壓后先認(rèn)路入口腳本與核心模塊的分工解壓后先別急著裝依賴對(duì)著文件結(jié)構(gòu)走一遍把每個(gè)文件在訓(xùn)練鏈路里的位置標(biāo)出來(lái)。這個(gè)包解壓后大致是下面這個(gè)結(jié)構(gòu)個(gè)別發(fā)布版本可能略有出入diffusion_tf/ ├── models/ │ ├── __init__.py # 模型包導(dǎo)出 │ └── nn.py # U-Net 噪聲預(yù)測(cè)網(wǎng)絡(luò) ├── diffusion_utils.py # 前向加噪、損失、采樣核心工具 ├── diffusion_utils_2.py # 兼容層/補(bǔ)充工具部分實(shí)現(xiàn)保留在這里 ├── tpu_utils.py # TPU 訓(xùn)練參數(shù)與分布式策略 ├── utils.py # 圖像保存、日志、EMA 等通用函數(shù) ├── scripts/ │ ├── run_cifar.py # CIFAR-10 訓(xùn)練與采樣入口 │ ├── run_celebahq.py # CelebA-HQ 訓(xùn)練與采樣入口 │ └── run_lsun.py # LSUN 訓(xùn)練與采樣入口 ├── requirements.txt # 依賴清單 ├── README.md # 使用說(shuō)明先讀它 └── samples.png # 參考生成結(jié)果用來(lái)和你的輸出對(duì)比models/nn.py 是 U-Net輸入帶噪圖 x_t 和時(shí)間步 t輸出與輸入同尺寸的預(yù)測(cè)噪聲。DDPM 用的 U-Net 比分割任務(wù)版本多了時(shí)間嵌入融合和更多殘差塊中間層還加了 attention參數(shù)規(guī)模通常在幾十百萬(wàn)量級(jí)。diffusion_utils.py 是核心工具集合訓(xùn)練時(shí)調(diào) q_sample 和 p_losses采樣時(shí)調(diào) p_sample后文調(diào)參主要待的地方也是這里。diffusion_utils_2.py 從命名習(xí)慣看像是保留函數(shù)集的兼容層如果主文件里某個(gè)函數(shù)在調(diào)用鏈上沒(méi)找到定義可以來(lái)這里翻一翻。tpu_utils.py 只在 TPU 訓(xùn)練時(shí)才會(huì)真正參與本機(jī)單卡訓(xùn)練可以直接忽略。utils.py 里通常裝著 EMA 更新和圖像保存邏輯EMA 權(quán)重對(duì)采樣質(zhì)量影響非常大后面避坑章會(huì)再提。scripts 目錄下三個(gè)入口腳本分別對(duì)應(yīng)三個(gè)數(shù)據(jù)集命令行參數(shù)不一樣跑之前用 --help 確認(rèn)。3.2 調(diào)用鏈一處訓(xùn)練step里nn.py與diffusion_utils.py如何配合訓(xùn)練腳本里一個(gè)標(biāo)準(zhǔn)的 step 只做三件事采樣時(shí)間步、加噪、預(yù)測(cè)噪聲求 loss。寫成偽代碼更直白# 訓(xùn)練循環(huán)中一次前向的典型順序 x_noisy q_sample(x_start, t, noise, alphas_cumprod) # 來(lái)自 diffusion_utils noise_pred model(x_noisy, t) # 來(lái)自 models/nn.py loss mean_squared_error(noise, noise_pred)這里的關(guān)鍵是模型完全不接收 x_0只接收 x_noisy 和 t。所以它對(duì)加噪程度的感知全部來(lái)自時(shí)間嵌入t 的編碼方式和 U-Net 里的特征融合方式直接決定模型能不能區(qū)分輕度噪聲和重度噪聲。改分辨率時(shí)U-Net 的輸入通道不變主要改中間特征維度和 attention 塊數(shù)改 T 時(shí)只需要改 beta 調(diào)度長(zhǎng)度和 alphas_cumprod網(wǎng)絡(luò)結(jié)構(gòu)完全不用動(dòng)。這個(gè)解耦設(shè)計(jì)是擴(kuò)散模型工程實(shí)現(xiàn)里最省心的地方。3.3 三個(gè)入口腳本run_cifar、run_celebahq、run_lsun分別適合什么場(chǎng)景腳本數(shù)據(jù)集圖像特點(diǎn)主要開(kāi)銷適合用途run_cifar.pyCIFAR-1032x32 小圖單卡可訓(xùn)練快速驗(yàn)證原理、課設(shè)首選run_celebahq.pyCelebA-HQ1024 高清人臉顯存敏感人臉生成、質(zhì)量評(píng)估run_lsun.pyLSUN類別眾多、量大數(shù)據(jù)讀取耗時(shí)長(zhǎng)類別條件生成實(shí)驗(yàn)run_cifar.py 是新手第一站小圖訓(xùn)練速度快跑一兩個(gè)小時(shí)就能看到像樣的輪廓。run_celebahq.py 和 run_lsun.py 在數(shù)據(jù)讀取上走不同目錄結(jié)構(gòu)CelebA-HQ 要求把圖片按文件夾整理好并生成文件列表LSUN 通常用 lmdb 格式腳本一般支持 --category 指定類別比如 bedroom、tower。跑之前先確認(rèn)數(shù)據(jù)目錄結(jié)構(gòu)是不是腳本預(yù)期的那樣這一條在避坑章會(huì)展開(kāi)。3.4 samples.png的正確用法不是裝飾圖是驗(yàn)收基準(zhǔn)包里那張 samples.png 不是給你看效果的裝飾圖而是驗(yàn)收基準(zhǔn)。正式訓(xùn)練前先打開(kāi)它記住正常生成結(jié)果該長(zhǎng)什么樣顏色分布自然、輪廓可辨認(rèn)、沒(méi)有大面積純色塊。訓(xùn)練過(guò)程中每隔一段時(shí)間跑一次采樣和 samples.png 并排對(duì)比。如果顏色分布接近說(shuō)明網(wǎng)絡(luò)學(xué)到了數(shù)據(jù)分布如果全是灰塊或高頻噪點(diǎn)按第 5 章的排查順序走別急著調(diào)網(wǎng)絡(luò)結(jié)構(gòu)。3.5 requirements.txt與README安裝和運(yùn)行前先看這兩個(gè)文件很多下載資源的 README 被當(dāng)成擺設(shè)這個(gè)包不建議這樣。README 里通常會(huì)寫數(shù)據(jù)集準(zhǔn)備的命令、版本要求、已知問(wèn)題比任何教程都貼近這個(gè)源碼包的實(shí)際狀態(tài)。requirements.txt 則鎖定了 TensorFlow、NumPy、Pillow 等依賴的版本范圍Python 環(huán)境里版本沖突是最常見(jiàn)的安裝失敗原因。先建虛擬環(huán)境再裝依賴能避開(kāi)和系統(tǒng) Python 包之間的沖突。4. 把環(huán)境搭起來(lái)并跑通CIFAR依賴安裝與參數(shù)選型4.1 環(huán)境準(zhǔn)備虛擬環(huán)境、依賴與幫助命令環(huán)境準(zhǔn)備這一步值得多花十分鐘后面能省下幾小時(shí)的排錯(cuò)時(shí)間。先把項(xiàng)目目錄拿過(guò)來(lái)建獨(dú)立虛擬環(huán)境裝依賴然后不要急著訓(xùn)練先看一遍幫助信息cd diffusion_tf python -m venv venv source venv/bin/activate pip install -r requirements.txt python scripts/run_cifar.py --helpPython 版本建議 3.8 及以上具體以 README 里寫的為準(zhǔn)。requirements.txt 里一般包含 TensorFlow、NumPy、Pillow 這些基礎(chǔ)庫(kù)直接 pip 安裝即可。--help 這一步是關(guān)鍵習(xí)慣不同發(fā)布版本的腳本參數(shù)名可能有出入比如有的叫 --data_dir有的叫 --data_root先確認(rèn)再訓(xùn)練能少踩很多坑。如果裝依賴時(shí)遇到 TensorFlow 版本和本機(jī) CUDA 不匹配優(yōu)先看 README 是否給了指定版本沒(méi)有的話選擇當(dāng)前 Python 版本下安裝最順利的官方版本。4.2 訓(xùn)練CIFAR入口命令與參數(shù)含義CIFAR-10 是 32x32 的小圖單張消費(fèi)級(jí)顯卡就能訓(xùn)練適合做第一個(gè)完整實(shí)驗(yàn)。訓(xùn)練命令按這個(gè)模板改python scripts/run_cifar.py \ --data_dir./data/cifar10 \ --out_dir./checkpoints \ --batch_size64 \ --num_epochs100 \ --lr2e-4腳本會(huì)先檢查數(shù)據(jù)目錄是否存在不存在時(shí)可能自動(dòng)下載 CIFAR-10也可能提示你手動(dòng)放置。參數(shù)含義按下面的表理解就夠參數(shù)典型值作用與注意事項(xiàng)data_dir./data/cifar10數(shù)據(jù)目錄目錄不存在時(shí)腳本一般會(huì)提示下載方式out_dir./checkpointscheckpoint 與日志輸出目錄建議按實(shí)驗(yàn)時(shí)間命名batch_size64單卡建議從 32 起步OOM 時(shí)優(yōu)先減小這一項(xiàng)num_epochs100DDPM 在小數(shù)據(jù)集上也需較多迭代看 loss 收斂曲線判斷l(xiāng)r2e-4原版常用 2e-4過(guò)高容易出現(xiàn) nan loss訓(xùn)練過(guò)程中注意看兩點(diǎn)loss 是否整體向下走以及日志里每個(gè) epoch 的速度是否穩(wěn)定。如果 loss 曲線平得像條直線先確認(rèn)數(shù)據(jù)加載是不是真的返回了圖像別在空數(shù)據(jù)上訓(xùn)練了一晚上。訓(xùn)練和采樣一般共用同一個(gè)入口腳本用 --mode 或單獨(dú)子命令區(qū)分具體以 --help 輸出的參數(shù)為準(zhǔn)。4.3 采樣與驗(yàn)收如何確認(rèn)模型真的學(xué)會(huì)了生成訓(xùn)練到一半或者全部結(jié)束后用采樣模式生成一批圖像python scripts/run_cifar.py \ --modesample \ --ckpt_dir./checkpoints \ --sample_dir./samples \ --num_samples64采樣命令通常不需要 data_dir因?yàn)橥评黼A段只從隨機(jī)噪聲出發(fā)不碰訓(xùn)練數(shù)據(jù)。腳本會(huì)從 checkpoint 恢復(fù)到指定 epoch然后從純?cè)肼暤?1000 步。檢查輸出時(shí)先看整體顏色分布再看邊緣銳利度顏色分布對(duì)了說(shuō)明網(wǎng)絡(luò)學(xué)到了數(shù)據(jù)分布邊緣清楚說(shuō)明高頻細(xì)節(jié)也學(xué)到了。如果生成結(jié)果全是同一種顏色或者所有圖長(zhǎng)得一模一樣基本是收斂異?；虿蓸庸絾?wèn)題去第 5 章對(duì)照排查。4.4 換數(shù)據(jù)集CelebA-HQ和LSUN的注意點(diǎn)跑通 CIFAR 之后再換數(shù)據(jù)集主要變化在分辨率和數(shù)據(jù)讀取。下面這張表是我在單卡環(huán)境下的典型配置參考數(shù)據(jù)集常用分辨率單卡顯存建議Batch Size 參考CIFAR-1032x324 GB 以上64CelebA-HQ128x128 或 256x25612 GB 以上16 或 8LSUN128x12812 GB 以上16CelebA-HQ 原圖是 1024x1024直接把原圖丟進(jìn)去訓(xùn)練普通消費(fèi)級(jí)顯卡根本頂不住。常見(jiàn)做法是先把圖片縮放到 128x128 或 256x256再做中心裁剪和歸一化run_celebahq.py 里的輸入管線通常已經(jīng)預(yù)留了 resize 配置你只需要把目標(biāo)分辨率參數(shù)填上。LSUN 的 bedroom 類有上百萬(wàn)張圖單機(jī)單卡全量訓(xùn)練要跑很久常見(jiàn)做法是先用子集或低分辨率做冒煙測(cè)試確認(rèn)數(shù)據(jù)和模型都正常再上全量。5. 避坑與常見(jiàn)問(wèn)題排查DDPM訓(xùn)練里最容易翻車的四個(gè)環(huán)節(jié)擴(kuò)散模型訓(xùn)練里的坑很多屬于看著像玄學(xué)其實(shí)都有跡可循的類型。我按自己踩過(guò)的順序整理四條每一條都是現(xiàn)象、原因、解決三步走。5.1 loss正常下降采樣結(jié)果卻全是噪聲或灰塊現(xiàn)象訓(xùn)練日志里 loss 曲線穩(wěn)定下滑看起來(lái)一切正常但跑出來(lái)的 samples 圖是一堆高頻噪點(diǎn)或者整片灰?guī)缀鯖](méi)有圖像結(jié)構(gòu)。原因通常有兩個(gè)。一是模型還在欠擬合階段CIFAR 上通常要跑比較多的迭代才能看到明顯輪廓訓(xùn)練幾十個(gè) epoch 就采樣結(jié)果當(dāng)然是噪聲。二是采樣代碼用了剛保存的原始權(quán)重而不是 EMA 權(quán)重。原始權(quán)重是訓(xùn)練過(guò)程的實(shí)時(shí)快照噪聲大EMA 權(quán)重是全程權(quán)重的滑動(dòng)平均推理時(shí)穩(wěn)得多。解決先確認(rèn) checkpoint 里保存的是否包含 EMA 參數(shù)如果 utils.py 里有 ema 模塊采樣前強(qiáng)制加載它。再檢查 p_sample 里最后一個(gè)時(shí)間步的 sigma 是否被置零如果沒(méi)置零最后一步的隨機(jī)噪聲會(huì)蓋掉整張圖。5.2 單卡GPU訓(xùn)練直接OOM顯存完全不夠現(xiàn)象訓(xùn)練腳本啟動(dòng)后沒(méi)跑幾個(gè) step直接報(bào)資源耗盡進(jìn)程退出。原因擴(kuò)散模型同時(shí)保存輸入圖和多個(gè)中間特征U-Net 的通道數(shù)、attention 層都比普通分類網(wǎng)絡(luò)吃顯存很多。很多人把 batch_size 設(shè)成分類任務(wù)的習(xí)慣值比如 128直接爆掉。解決先把 batch_size 降到 16 甚至 8確認(rèn)能跑起來(lái)再往上加。如果還想提升訓(xùn)練效率打開(kāi)混合精度在保留 float32 主權(quán)重的同時(shí)用 float16 計(jì)算腳本沒(méi)有內(nèi)置的話用梯度累積模擬大 batch也就是多個(gè)小 batch 前向求平均后再反向更新。OOM 報(bào)錯(cuò)信息里會(huì)給出當(dāng)前申請(qǐng)到的顯存大小拿它和顯卡總顯存對(duì)比能粗略估算出下一步該調(diào)到多少。5.3 CelebA-HQ與LSUN數(shù)據(jù)加載經(jīng)常報(bào)路徑或尺寸錯(cuò)誤現(xiàn)象換到 CelebA-HQ 或 LSUN 后訓(xùn)練腳本報(bào)文件找不到、lmdb 讀取失敗或者圖片尺寸不一致導(dǎo)致 reshape 出錯(cuò)。原因這兩個(gè)數(shù)據(jù)集不是直接解壓就能用的。CelebA-HQ 需要按目錄整理圖片并準(zhǔn)備文件列表LSUN 是 lmdb 格式且分多個(gè)類別腳本預(yù)期的目錄結(jié)構(gòu)和網(wǎng)上直接下載的原始包經(jīng)常對(duì)不上。解決先按 README 或腳本頂部注釋要求的目錄結(jié)構(gòu)放置數(shù)據(jù)。我一般會(huì)在正式訓(xùn)練前寫一個(gè)幾行的數(shù)據(jù)自檢腳本打印一個(gè) batch 的 shape 和像素范圍確認(rèn)沒(méi)問(wèn)題再跑。批量 resize 到統(tǒng)一分辨率再存成 tfrecord 或 numpy 格式能省后面很多事。5.4 訓(xùn)練中后期頻繁出現(xiàn)nan loss模型直接崩掉現(xiàn)象loss 正常降了幾千步突然變成 nan后續(xù) step 全部異常checkpoint 也無(wú)法繼續(xù)加載。原因?qū)W習(xí)率偏高在深層網(wǎng)絡(luò)里會(huì)放大梯度diffusion 模型尤其敏感混合精度下 float16 的 loss scaling 不夠分母里出現(xiàn)極小值會(huì)溢出U-Net 某層歸一化的 epsilon 默認(rèn)值太小也會(huì)導(dǎo)致除零。解決把學(xué)習(xí)率降到 1e-4 或 5e-5先確認(rèn)學(xué)習(xí)率調(diào)度沒(méi)問(wèn)題如果用混合精度把 loss scaling 初始值調(diào)大并在反向更新前跳過(guò) inf 梯度在 nn.py 和 diffusion_utils.py 里把歸一化層的 epsilon 從 1e-5 改成 1e-6 再試。大多數(shù) nan 問(wèn)題都能被這三步定位到。改完先重跑一小段確認(rèn) loss 能跨過(guò)之前崩潰的步數(shù)閾值。5.5 一類通用排查習(xí)慣先縮小問(wèn)題空間再上全量訓(xùn)練如果上面的現(xiàn)象都對(duì)不上我習(xí)慣先做一個(gè)極端小的冒煙測(cè)試把圖像縮到 16x16 或 32x32batch_size 設(shè)為 4跑 20 個(gè) epoch看模型能不能在這么小的數(shù)據(jù)上過(guò)擬合。這個(gè)測(cè)試能一次性驗(yàn)證數(shù)據(jù)管線、模型前向、loss 計(jì)算、checkpoint 保存四條鏈路是否正常。如果小規(guī)模都跑不通問(wèn)題一定在代碼層面如果小規(guī)模能過(guò)擬合全量跑不通再往顯存、數(shù)據(jù)量、學(xué)習(xí)率方向查。這套檢查順序能篩掉大半無(wú)聲的配置錯(cuò)誤比直接上全量訓(xùn)練然后盯著日志猜原因高效得多。6. 進(jìn)階技巧給DDPM接入DDIM采樣與類別條件控制6.1 把DDPM采樣換成DDIM用50步替代1000步DDPM 采樣要迭代 1000 步生產(chǎn)環(huán)境里太慢。DDIM 的核心改動(dòng)是把隨機(jī)項(xiàng)系數(shù)設(shè)為可調(diào)當(dāng) eta0 時(shí)采樣過(guò)程完全確定性50 步左右就能生成質(zhì)量可接受的圖像。代碼上只改采樣函數(shù)一小段訓(xùn)練完全不動(dòng)# DDIM單步更新eta0 時(shí)完全確定可用 50 步左右生成 def ddim_sample(model, x_t, t, acp_t, acp_prev, eta0.0): eps model(x_t, t) x0_pred (x_t - torch.sqrt(1 - acp_t) * eps) / torch.sqrt(acp_t) sigma eta * torch.sqrt((1 - acp_prev) / (1 - acp_t) * (1 - acp_t / acp_prev)) noise torch.randn_like(x_t) if eta 0 else 0 x_prev torch.sqrt(acp_prev) * x0_pred \ torch.sqrt(1 - acp_prev - sigma**2) * eps sigma * noise return x_prev在 diffusion_utils.py 里另寫一個(gè) ddim_sample 函數(shù)采樣循環(huán)只改調(diào)用訓(xùn)練代碼一行都不用動(dòng)。eta 調(diào)大能增加生成多樣性但也會(huì)引入隨機(jī)顆粒感常見(jiàn)做法是先設(shè) 0確認(rèn)效果后再微調(diào)。6.2 給U-Net加類別條件如果你想做按指定標(biāo)簽生成常見(jiàn)做法是把類別標(biāo)簽做 embedding然后和 time embedding 相加后一起注入殘差塊。用 run_lsun.py 做類別實(shí)驗(yàn)時(shí)這個(gè)入口本身就暗示了類別支持的路徑。改動(dòng)集中在一個(gè)位置在模型 forward 里加 class_embedding維度對(duì)齊 time embedding后續(xù)融合方式照抄時(shí)間嵌入的融合邏輯。要注意類別數(shù)變化時(shí) embedding 表也要同步重建別沿用舊 checkpoint。6.3 EMA權(quán)重的保存習(xí)慣生成質(zhì)量最穩(wěn)的經(jīng)驗(yàn)是訓(xùn)練期維護(hù)一份 EMA 權(quán)重checkpoint 主權(quán)重每 N 步保存一次EMA 權(quán)重單獨(dú)保存評(píng)估或者采樣本時(shí)全部用 EMA。這個(gè)習(xí)慣在 CIFAR 上可能區(qū)別不大在 CelebA-HQ 這種高清數(shù)據(jù)集上差別非常明顯。從那以后我每接手一個(gè) diffusion 代碼包都強(qiáng)制先跑通一張圖的過(guò)擬合冒煙測(cè)試再改 batch 和分辨率上全量訓(xùn)練。這個(gè)習(xí)慣幫我避掉了無(wú)數(shù)次無(wú)聲的配置錯(cuò)誤。希望幫到你。本文還有配套的精品資源點(diǎn)擊獲取