07 - DiT与Transformer扩散架构¶

import torch
import torch.nn as nn

class PatchEmbed(nn.Module):  # 继承nn.Module定义网络层
    """将潜空间特征图转换为patch序列"""

    def __init__(self, img_size=32, patch_size=2, in_channels=4, embed_dim=768):
        super().__init__()  # super()调用父类方法
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # x: [B, C, H, W] → [B, D, H/p, W/p] → [B, N, D]
        x = self.proj(x)                    # [B, D, H/p, W/p]
        x = x.flatten(2).transpose(1, 2)    # [B, N, D]
        return x

class AdaLNZero(nn.Module):
    """Adaptive Layer Norm Zero — DiT的核心条件注入模块"""

    def __init__(self, dim, cond_dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
        # 预测 gamma, beta, alpha 三组参数
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 6 * dim)  # 为attention和MLP各预测3个参数
        )
        # 零初始化
        nn.init.zeros_(self.adaLN_modulation[-1].weight)  # [-1]负索引取最后元素
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, x, cond):
        # cond: [B, cond_dim] → 预测6组调制参数
        (gamma1, beta1, alpha1,
         gamma2, beta2, alpha2) = self.adaLN_modulation(cond).chunk(6, dim=-1)

        return gamma1, beta1, alpha1, gamma2, beta2, alpha2

    def modulate(self, x, gamma, beta):
        return self.norm(x) * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)  # unsqueeze增加一个维度

import torch
import torch.nn as nn
import math

class SinusoidalPosEmb(nn.Module):
    """正弦时间步编码"""
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None].float() * emb[None, :]
        return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)  # torch.cat沿已有维度拼接张量

class DiTBlock(nn.Module):
    """DiT Transformer Block with AdaLN-Zero"""

    def __init__(self, dim, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        # Self-Attention
        self.qkv = nn.Linear(dim, 3 * dim)
        self.proj = nn.Linear(dim, dim)

        # FFN
        mlp_hidden = int(dim * mlp_ratio)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_hidden),
            nn.GELU(),
            nn.Linear(mlp_hidden, dim),
        )

        # AdaLN-Zero: 两层各需 gamma, beta, alpha → 6 * dim
        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
        self.adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, 6 * dim),
        )
        # 关键：零初始化
        nn.init.zeros_(self.adaLN[-1].weight)
        nn.init.zeros_(self.adaLN[-1].bias)

    def forward(self, x, cond):
        """
        x: [B, N, D] — patch token序列
        cond: [B, D] — 条件嵌入(时间步+类别)
        """
        B, N, D = x.shape

        # 预测6组调制参数
        gamma1, beta1, alpha1, gamma2, beta2, alpha2 = \
            self.adaLN(cond).chunk(6, dim=-1)

        # --- Attention分支 ---
        h = self.norm1(x) * (1 + gamma1.unsqueeze(1)) + beta1.unsqueeze(1)
        qkv = self.qkv(h).reshape(B, N, 3, self.num_heads, D // self.num_heads)  # 重塑张量形状
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        h = (attn @ v).transpose(1, 2).reshape(B, N, D)
        h = self.proj(h)
        x = x + alpha1.unsqueeze(1) * h  # 零初始化门控

        # --- FFN分支 ---
        h = self.norm2(x) * (1 + gamma2.unsqueeze(1)) + beta2.unsqueeze(1)
        h = self.mlp(h)
        x = x + alpha2.unsqueeze(1) * h  # 零初始化门控

        return x

class DiT(nn.Module):
    """简化版 Diffusion Transformer"""

    def __init__(
        self,
        img_size=32,
        patch_size=2,
        in_channels=4,
        dim=768,
        depth=12,
        num_heads=12,
        num_classes=10,
        mlp_ratio=4.0,
    ):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2

        # Patch Embedding
        self.patch_embed = nn.Conv2d(
            in_channels, dim, kernel_size=patch_size, stride=patch_size
        )
        # 位置编码 (固定sin-cos)
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.num_patches, dim)
        )

        # 条件编码
        self.time_embed = nn.Sequential(
            SinusoidalPosEmb(dim),
            nn.Linear(dim, dim),
            nn.GELU(),
            nn.Linear(dim, dim),
        )
        self.class_embed = nn.Embedding(num_classes, dim)

        # Transformer Blocks
        self.blocks = nn.ModuleList([
            DiTBlock(dim, num_heads, mlp_ratio) for _ in range(depth)
        ])

        # 输出层
        self.final_norm = nn.LayerNorm(dim, elementwise_affine=False)
        self.final_adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim, 2 * dim),
        )
        self.final_proj = nn.Linear(dim, patch_size ** 2 * in_channels * 2)
        # 2倍channels用于预测噪声和对角方差

        self._init_weights()

    def _init_weights(self):
        # 初始化位置编码
        nn.init.normal_(self.pos_embed, std=0.02)
        # 零初始化输出层
        nn.init.zeros_(self.final_proj.weight)
        nn.init.zeros_(self.final_proj.bias)
        nn.init.zeros_(self.final_adaLN[-1].weight)
        nn.init.zeros_(self.final_adaLN[-1].bias)

    def unpatchify(self, x, H, W):
        """[B, N, p*p*C] → [B, C, H, W]"""
        p = self.patch_size
        c = x.shape[-1] // (p * p)
        h, w = H // p, W // p
        x = x.reshape(-1, h, w, p, p, c)
        x = x.permute(0, 5, 1, 3, 2, 4).reshape(-1, c, H, W)
        return x

    def forward(self, x, t, y):
        """
        x: [B, C, H, W] — 带噪声的潜变量
        t: [B] — 时间步
        y: [B] — 类别标签
        """
        B, C, H, W = x.shape

        # Patchify + 位置编码
        x = self.patch_embed(x).flatten(2).transpose(1, 2)  # [B, N, D]
        x = x + self.pos_embed

        # 条件编码
        cond = self.time_embed(t) + self.class_embed(y)     # [B, D]

        # Transformer Blocks
        for block in self.blocks:
            x = block(x, cond)

        # 输出投影
        gamma, beta = self.final_adaLN(cond).chunk(2, dim=-1)
        x = self.final_norm(x) * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)
        x = self.final_proj(x)   # [B, N, p*p*2C]
        x = self.unpatchify(x, H, W)  # [B, 2C, H, W]

        noise_pred, var_pred = x.chunk(2, dim=1)
        return noise_pred, var_pred

# === 使用示例 ===
if __name__ == "__main__":
    model = DiT(
        img_size=32, patch_size=2, in_channels=4,
        dim=384, depth=12, num_heads=6, num_classes=10
    )
    print(f"DiT-S/2 参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")

    x = torch.randn(2, 4, 32, 32)  # 潜空间输入
    t = torch.randint(0, 1000, (2,))
    y = torch.randint(0, 10, (2,))

    noise_pred, var_pred = model(x, t, y)
    print(f"输出形状: noise={noise_pred.shape}, var={var_pred.shape}")
    # 输出: noise=torch.Size([2, 4, 32, 32]), var=torch.Size([2, 4, 32, 32])