Transformer、Stable Diffusion 与 LLM 模型对比及实现

接下来我们看看具体的代码实现。这里以 DDPM（Denoising Diffusion Probabilistic Models）为例，展示如何配置训练参数并启动流程。

主程序入口

这个脚本负责配置所有训练或评估所需的超参数，并根据模式调用相应的函数。

# DDPM 主程序入口
# 使用方法：Python Main.py
# 使用默认配置训练或评估
from Diffusion.Train import train, eval

def main(model_config=None):
    # 主函数：配置参数并启动训练或评估
    # args: model_config 可选的配置字典，如果提供则覆盖默认配置
    
    # 模式配置字典：包含所有训练和评估所需的超参数
    modelConfig = {
        # 运行模式：train 训练模式，从头开始训练或继续训练模型；eval 评估模式，加载已训练模型并生成图像
        "state": "train",
        
        # 训练超参数
        "epoch": 200,           # 训练的总轮数
        "batch_size": 80,       # 每个批次的样本数量
        "lr": 1e-4,             # 初始学习率
        "multiplier": 2,        # 学习率预热后的倍数
        "grad_clip": 1,         # 梯度裁剪阈值，防止梯度爆炸
        
        # 扩散过程参数
        "T": 1000,              # 扩散步数前向过程的步数
        "beta_1": 1e-4,         # 初始噪声系数
        "beta_T": 0.02,         # 最终噪声系数
        
        # UNet 模型架构参数
        "channel": 128,         # 基础通道数
        "channel_mult": [1, 2, 3, 4], # 通道倍数列表
        "attn": [2],            # 在哪些层级添加注意力机制
        "num_res_blocks": 2,    # 每个分辨率级别的残差块数量
        "dropout": 0.15,        # Dropout 比率
        
        # 数组相关参数
        "img_size": 32,         # 图像尺寸
        
        # 设备配置
        "device": "cuda:0",     # 计算设备
        
        # 模型权重相关
        "training_load_weight": None, # 训练时加载的预训练权重文件名
        "save_weight_dir": "./Checkpoints/", # 模型权重保存目录
        "test_load_weight": "ckpt_199_.pt", # 评估时加载的模型权重文件名
        
        # 图像保存相关
        "sampled_dir": "./SampledImgs/", # 生成图像的保存目录
        "sampledNoisyImgName": "NoisyNoGuidenceImgs.png", # 初始噪声图像文件名
        "sampledImgName": "SampledNoGuidenceImgs.png", # 生成图像文件名
        "nrow": 8               # 保存图像时每行显示的图像数量
    }
    
    # 如果提供了自定义配置，则使用自定义配置覆盖默认值
    if model_config is not None:
        modelConfig.update(model_config)
    
    # 根据配置选择运行模式
    if modelConfig["state"] == "train":
        train(modelConfig)
    else:
        eval(modelConfig)

if __name__ == '__main__':
    main()

训练与评估代码解析

训练模块封装了 DDPM 的前向扩散过程和反向去噪过程。核心思想是通过神经网络学习从噪声中恢复原始图像。

训练流程

训练的目标是最小化预测噪声和真实噪声之间的 MSE 损失。以下是训练函数的关键逻辑：

import os
from typing import Dict
import torch
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR10
from torchvision.utils import save_image
from Diffusion import GaussianDiffusionSampler, GaussianDiffusionTrainer
from Diffusion.Model import UNet
from Scheduler import GradualWarmupScheduler

def train(modelConfig: Dict):
    # 设置计算设备，GPU 或 CPU
    device = torch.device(modelConfig["device"])
    
    # 数据集准备：加载 CIFAR-10 训练集
    dataset = CIFAR10(
        root='./CIFAR10', 
        train=True, 
        download=True, 
        transform=transforms.Compose([
            transforms.RandomHorizontalFlip(), 
            transforms.ToTensor(), 
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
    )
    
    # 创建数据加载器
    dataloader = DataLoader(
        dataset, 
        batch_size=modelConfig["batch_size"], 
        shuffle=True, 
        num_workers=4, 
        drop_last=True, 
        pin_memory=True
    )
    
    # 模型初始化：创建 UNet 模型
    net_model = UNet(
        T=modelConfig["T"], 
        ch=modelConfig["channel"], 
        ch_mult=modelConfig["channel_mult"], 
        attn=modelConfig["attn"], 
        num_res_blocks=modelConfig["num_res_blocks"], 
        dropout=modelConfig["dropout"]
    ).to(device)
    
    # 如果指定了预训练权重，则加载
    if modelConfig["training_load_weight"] is not None:
        net_model.load_state_dict(torch.load(os.path.join(modelConfig["save_weight_dir"], modelConfig["training_load_weight"]), map_location=device))
    
    # 优化器设置：使用 AdamW 优化器
    optimizer = torch.optim.AdamW(net_model.parameters(), lr=modelConfig["lr"], weight_decay=1e-4)
    
    # 学习率调度器设置：余弦退火 + 预热
    cosineScheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=modelConfig["epoch"], eta_min=0)
    warmUpScheduler = GradualWarmupScheduler(
        optimizer=optimizer, 
        multiplier=modelConfig["multiplier"], 
        warm_epoch=modelConfig["epoch"] // 10, 
        after_scheduler=cosineScheduler
    )
    
    # 扩散训练器初始化
    trainer = GaussianDiffusionTrainer(
        net_model, 
        modelConfig["beta_1"], 
        modelConfig["beta_T"], 
        modelConfig["T"]
    ).to(device)
    
    # 开始训练循环
    for e in range(modelConfig["epoch"]):
        with tqdm(dataloader, dynamic_ncols=True) as tqdmDataLoader:
            for images, labels in tqdmDataLoader:
                optimizer.zero_grad()
                x_0 = images.to(device)
                
                # 前向传播，计算损失
                loss = trainer(x_0).sum() / 1000.0
                
                # 反向传播
                loss.backward()
                
                # 梯度裁剪
                torch.nn.utils.clip_grad_norm_(net_model.parameters(), modelConfig["grad_clip"])
                
                # 更新模型参数
                optimizer.step()
                
                # 更新进度条显示信息
                tqdmDataLoader.set_postfix({"epoch": e, "loss": loss.item(), "img shape": x_0.shape, "LR": optimizer.state_dict()['param_groups'][0]["lr"]})
        
        # 每个 epoch 结束后更新学习率
        warmUpScheduler.step()
        
        # 保存模型检查点
        torch.save(net_model.state_dict(), os.path.join(modelConfig["save_weight_dir"], 'ckpt_' + str(e) + ".pt"))

评估与采样

评估阶段实际上是生成新图像的过程。我们从标准高斯分布采样初始噪声，然后逐步去噪。

def eval(modelConfig: Dict):
    with torch.no_grad():
        device = torch.device(modelConfig["device"])
        
        # 加载训练好的模型
        model = UNet(
            T=modelConfig["T"], 
            ch=modelConfig["channel"], 
            ch_mult=modelConfig["channel_mult"], 
            attn=modelConfig["attn"], 
            num_res_blocks=modelConfig["num_res_blocks"], 
            dropout=0
        ).to(device)
        
        ckpt = torch.load(os.path.join(modelConfig["save_weight_dir"], modelConfig["test_load_weight"]), map_location=device)
        model.load_state_dict(ckpt)
        print("model load weight done.")
        
        model.eval()
        
        # 初始化采样器
        sampler = GaussianDiffusionSampler(
            model, 
            modelConfig["beta_1"], 
            modelConfig["beta_T"], 
            modelConfig["T"]
        ).to(device)
        
        # 从标准高斯分布采样初始噪声
        noisyImage = torch.randn(size=[modelConfig["batch_size"], 3, 32, 32], device=device)
        
        # 保存初始噪声图像
        saveNoisy = torch.clamp(noisyImage * 0.5 + 0.5, 0, 1)
        save_image(saveNoisy, os.path.join(modelConfig["sampled_dir"], modelConfig["sampledNoisyImgName"]), nrow=modelConfig["nrow"])
        
        # 执行采样过程
        print("Starting sampling process...")
        sampledImgs = sampler(noisyImage)
        
        # 将生成的图像转换到 [0, 1] 范围
        sampledImgs = sampledImgs * 0.5 + 0.5
        
        # 保存生成的图像
        save_image(sampledImgs, os.path.join(modelConfig["sampled_dir"], modelConfig["sampledImgName"]), nrow=modelConfig["nrow"])
        print(f"Sampling completed! Images saved to {modelConfig['sampled_dir']}")

通过这段代码，我们可以清晰地看到从数据加载、模型构建、训练循环到最终采样的完整链路。实际运行时，建议确保 GPU 资源充足，因为 DDPM 训练对计算要求较高。