Transformer、LLM 与 Stable Diffusion 模型对比分析

开源的影响：对于 LLM 和 Stable Diffusion，开源极大地降低了应用门槛，促进了安全性研究、垂直领域适配、成本下降和技术民主化。

三、Stable Diffusion 模型实现示例

# DDPM 主程序入口
# 这个文件是 DDPM 模型的启动脚本，负责
# 1. 配置所有训练/评估参数
# 2. 根据配置选择训练或评估模式
# 3. 调用相应的训练或评估函数
# 使用方法：Python Main.py
# 使用默认配置训练或评估
from Diffusion.Train import train, eval

def main(model_config=None):
    # 主函数：配置参数并启动训练或评估
    # args: model_config 可选的配置字典，如果提供则覆盖默认配置
    
    # 模式配置字典：包含所有训练和评估所需的超参数
    modelConfig = {
        # 运行模式
        "state": "train",  # train 或 eval
        # train 训练模式，从头开始训练或继续训练模型
        # eval 评估模式，加载已训练模型并生成图像
        
        # 训练超参数
        "epoch": 200,          # 训练的总轮数
        "batch_size": 80,      # 每个批次的样本数量
        "lr": 1e-4,            # 初始学习率，learning rate
        "multiplier": 2,       # 学习率预热后的倍数，预热后学习率 = lr * multiplier
        "grad_clip": 1,        # 梯度裁剪阈值，防止梯度爆炸
        
        # 扩散过程参数
        "T": 1000,             # 扩散步数前向过程的步数
        # 扩散过程 x_0 -> x_1 -> ... -> x_T
        # 更多步数通常能生成更高质量的图像，但计算成本更高
        "beta_1": 1e-4,        # 初始噪声系数（第 1 步的噪声方差）
        "beta_T": 0.02,        # 最终噪声系数（第 T 步的噪声方差）
        # beta_t 定义了每一步添加的噪声量
        # 通常从很小的值 1e-4 线性增加到较大的值 0.02
        # 这确保了前向过程能逐渐将图像变成纯噪声
        
        # UNet 模型架构参数
        "channel": 128,        # 基础通道数，第一层的通道数
        "channel_mult": [1, 2, 3, 4],  # 通道倍数列表
        # 定义了下采样过程中通道数的变化
        # 例如：[1, 2, 3, 4] 表示：
        # - 第 1 层：128 通道
        # - 第 2 层：128*2 = 256 通道
        # - 第 3 层：128*3 = 384 通道
        # - 第 4 层：128*4 = 512 通道
        "attn": [2],           # 在哪些层级添加注意力机制
        # 例如，[2] 表示在第 2 个分辨率级别添加注意力层
        # 注意力机制能帮助模型捕获长距离依赖，提升生成质量
        "num_res_blocks": 2,   # 每个分辨率级别的残差块数量
        # 残差块是 UNet 的基本构建单元，包含时间步嵌入，time embedding，卷积层，可选的注意力层
        "dropout": 0.15        # Dropout 比率，训练时的正则化
        # Dropout 随即将部分神经元输出置零，防止过拟合
        
        # 数组相关参数
        "img_size": 32,        # 图像尺寸
        
        # 设备配置
        "device": "cuda:0",    # 计算设备
        # cuda:0 使用第 1 块 GPU，如果有 CPU 使用 CPU 训练会很慢
        # 注意 DDPM 训练需要大量计算，强烈建议使用 GPU
        
        # 模型权重相关
        "training_load_weight": None,  # 训练时加载的预训练权重文件名（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() 函数
    main()

四、训练代码

# DDPM (Denoising Diffusion Probabilistic) 训练和评估模块
# DDPM 核心思想
# 前向过程 Forward Process 逐步向真实图像添加高斯噪声，直到变成纯噪声
# q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1}, beta_t * I)
# 经过 T 步后 x_t 近似为标准高斯分布 N(0, I)
# 反向过程 Reverse Process 训练神经网络学习从噪声逐步去噪，恢复原始图像
# p_theta(x_{t-1} | x_t) = N(x_{t-1}; mu_theta(x_t, t), sigma_t^2 * I)
# 模型预测每一步的噪声，然后从 x_t 中减去预测的噪声得到 x_{t-1}
# 训练目标：最小化预测噪声和真实噪声之间的 MSE 损失
# L = E[||epsilon - epsilon_theta(x_t, t)||^2]
# 其中，epsilon 是前向过程中添加的真实噪声，epsilon_theta 是模型预测的噪声

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):
    # DDPM 模型训练函数
    # 训练流程
    # 1. 加载数据集 CIFAR-10
    # 2. 初始化 UNet 模型用于预测噪声
    # 3. 设置优化器和学习率调度器
    # 4. 使用 GaussianDiffusionTrainer 进行训练循环
    # 5. 随机采样时间步，根据 t 向图像添加噪声得到 x_t
    # 6. 模型预测噪声，计算损失并反向传播
    
    # args: modelConfig: 包含所有训练配置参数的字典
    
    # 设置计算设备，GPU 或 CPU
    device = torch.device(modelConfig["device"])
    
    # 数据集准备：加载 CIFAR-10 训练集
    # CIFAR-10 10 类彩色图像，每张图像 32x32 像素
    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)),  # 归一化到 [-1, 1] 范围，因为 DDPM 通常在这个范围内工作
        ])
    )
    
    # 创建数据加载器
    dataloader = DataLoader(
        dataset,
        batch_size=modelConfig["batch_size"],  # 批次大小
        shuffle=True,                          # 每个 epoch 打乱数据
        num_workers=4,                         # 数据加载的并行进程数
        drop_last=True,                        # 丢弃最后一个不完整的批次
        pin_memory=True                        # 将数据固定在内存中，加速 GPU 传输
    )
    
    # 模型初始化：创建 UNet 模型
    # UNet 是 DDPM 的核心网络，用于预测每一步噪声
    # 输入：带噪声的图像 x_t 和时间步 t
    # 输出：预测的噪声 epsilon_theta(x_t, t)
    net_model = UNet(
        T=modelConfig["T"],            # 扩散步数通常为 1000
        ch=modelConfig["channel"],     # 基础通道数
        ch_mult=modelConfig["channel_mult"],  # 通道倍数列表，如 [1,2,3,4] 表示下采样时通道数逐渐增加
        attn=modelConfig["attn"],      # 在哪些层级添加注意力机制
        num_res_blocks=modelConfig["num_res_blocks"],  # 每个分辨率级别的残差块数量
        dropout=modelConfig["dropout"] # 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, last_epoch=-1)
    
    # 预热调度器，在训练初期逐渐增加学习率，然后切换到余弦退火，这有助于训练初期的稳定性
    warmUpScheduler = GradualWarmupScheduler(optimizer=optimizer, multiplier=modelConfig["multiplier"], warm_epoch=modelConfig["epoch"]//10, after_scheduler=cosineScheduler)
    
    # 扩散训练器初始化：GaussianDiffusionTrainer 封装了 DDPM 的前向扩散过程
    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()
                
                # 将图像移到指定设备，GPU
                x_0 = images.to(device)
                # x_0 是原始干净图像，形状：batch_size, 3, 32, 32
                
                # 前向传播，计算损失
                # trainer.x_0 内部会：
                # 1. 随机采样时间步 t
                # 2. 采样噪声 epsilon
                # 3. 计算 x_t = sqrt(alpha_bar_t)* x_0 + sqrt(1 - alpha_bar_t) * epsilon
                # 4. 模型预测噪声 epsilon_theta = model(x_t, t)
                # 5. 计算 MSE 损失 loss = ||epsilon - epsilon_theta||^2
                # 返回每个样本的损失，形状：batch_size
                loss = trainer(x_0).sum() / 1000.0  # 求和后除以 1000 进行缩放数值稳定性
                
                # 反向传播，计算梯度
                loss.backward()
                
                # 梯度裁剪，防止梯度爆炸：将所有参数的梯度范围裁剪到 grad_clip 以内
                torch.nn.utils.clip_grad_norm_(net_model.parameters(), modelConfig["grad_clip"])
                
                # 更新模型参数
                optimizer.step()
                
                # 更新进度条显示信息
                tqdmDataLoader.set_postfix(epoch=e, loss=loss.item())
        
        # 每个 epoch 结束后更新学习率
        warmUpScheduler.step()
        
        # 保存模型检查点：每个 epoch 都保存，方便后续选择最大模型或继续训练
        torch.save(net_model.state_dict(), os.path.join(modelConfig["save_weight_dir"], 'ckpt_' + str(e) + '.pt'))

def eval(modelConfig: Dict):
    # DDPM 模型评估/采样函数
    # 评估流程，实际上是生成新图像
    # 1. 加载训练好的模型
    # 2. 从标准高斯分布采样初始噪声 x_T
    # 3. 使用 GaussianDiffusionSampler 逐步去噪，从 t=T 到 t=0，逐步预测并去除噪声，最终得到生成的图像 x_0
    
    # args: modelConfig 包含所有评估配置参数的字典
    
    # 使用 torch.no_grad() 禁用梯度计算，节省内存和加速
    with torch.no_grad():
        device = torch.device(modelConfig["device"])
        
        # 加载训练好的模型：创建模型结构，评估时 dropout 设为 0
        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  # 评估时不需要 dropout
        ).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.")
        
        # 设置为评估模式，影响 BatchNorm, Dropout 等层的行为
        model.eval()
        
        # 初始化采样器：GaussianDiffusionSampler 封装了 DDPM 的反向去噪过程
        sampler = GaussianDiffusionSampler(
            model,
            modelConfig["beta_1"],
            modelConfig["beta_T"],
            modelConfig["T"]
        ).to(device)
        
        # ========== 生成图像 ==========
        # 从标准高斯分布采样初始噪声：这是反向过程的起点 x_T ~ N(0, I)
        noisyImage = torch.randn(size=[modelConfig["batch_size"], 3, 32, 32], device=device)
        
        # 保存初始噪声图像（用于可视化对比）：将 [-1, 1] 范围转换回 [0, 1] 用于保存
        saveNoisy = torch.clamp(noisyImage * 0.5 + 0.5, 0, 1)
        save_image(saveNoisy, os.path.join(modelConfig["sampled_dir"], modelConfig["sampledNoisyImgName"]), nrow=modelConfig["nrow"])
        
        # 执行采样过程：从噪声逐步去噪生成图像
        # sampler 内部会执行 T 步去噪：for t in [T, T-1, ..., 1]: 预测噪声 epsilon_theta = model(x_t, t)，计算 x_{t-1} 的均值，采样 x_{t-1} ~ N(mean, var)，最终得到 x_0
        print("Starting sampling process...")
        sampledImgs = sampler(noisyImage)
        
        # 将生成的图像从 [-1, 1] 范围转换到 [0, 1] 范围，便于保存和可视化
        sampledImgs = sampledImgs * 0.5 + 0.5  # [0 ~ 1]
        
        # 保存生成的图像
        save_image(sampledImgs, os.path.join(modelConfig["sampled_dir"], modelConfig["sampledImgName"]), nrow=modelConfig["nrow"])
        print(f"Sampling completed! Images saved to {modelConfig['sampled_dir']}")