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