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