企业官网建设流程全解析

1. 从扩散模型到DiT：骨干网络的一次模式切换

近两年生成模型圈子里最明显的一个变化，就是扩散模型的骨干网络从U-Net逐渐迁移到Transformer架构。DiT（Diffusion Transformers）就是这个趋势里最有代表性的一支。很多刚接触这个方向的人会有一个疑问：U-Net明明已经很成熟了，为什么还要费劲去换成Transformer？这篇文章就围绕DiT的原理和代码实现展开，带你把它从里到外拆一遍，看完你不仅知道它在干什么，还能直接照着代码上手搭一个能用的版本。

适合谁看？两类人最合适：一是已经写过简单DDPM、现在想进一步理解“工程化扩散模型”该如何设计架构的人；二是被各种模型名绕晕，想系统搞懂DiT、adaLN、Patchify这些概念到底对应什么模块、解决什么问题的新手。我会尽量把公式掰开讲，代码也会给全。

1.1 先把扩散模型的老底翻出来

扩散模型不是什么魔法，它的思路是反着来的。前向过程就是往一张干净图片上不断加高斯噪声：第1步加一点，第2步加得更多，走到几百步以后，图片基本成了一团什么信息都没有的纯噪声。反向过程则是让神经网络学会“倒放”，从这团噪声里一步步把原图恢复出来。

用数字表达前向过程就是：

q(x_t | x_0) = N(x_t; sqrt(α̅_t) * x_0, (1 - α̅_t) * I)

意思是在时间步t，当前样本可以看作由原图x_0乘以一个衰减系数，再加上σ级别的噪声合成。训练目标是让网络预测某个被刻意加进去的噪声ε：

L = E_{t,x0,ε} [ || ε_θ(x_t, t) - ε ||^2 ]

这是我个人觉得最好用的一个理解方式：扩散模型的每一步其实都是在教网络“估计噪声”。一旦噪声估计得准，采样时不停减掉噪声就行。类比一下，往清水里滴一滴墨水，墨迹会慢慢扩散到整杯水；训练时你记录了墨水原本那一下的位置，采样时就是倒过来做——从整杯均匀的墨水色里，一点点把“墨水抽走”。DiT在这个架构里干的活，就是那个负责“怎么抽墨水”的智能去噪器。

1.2 U-Net做得好好的，为什么非要换掉

U-Net的成名得益于它对图像的强归纳偏置：卷积天然擅长提取局部纹理，编码器配解码器加跳连接，让细节信息有“直通车”可以跨层传递。这套设计在低分辨率、小规模生成任务上又稳又快。可一旦任务规模变大，它的短板就藏不住了。

第一，U-Net的下采样/上采样结构对空间分辨率有明确设计假设，当你想直接增大生成图片尺寸时，往往需要改动网络拓扑，不能简单“堆参数”。第二，U-Net内部的attention机制通常只在瓶颈层启用，在模型容量和全局覆盖范围上天然受限。第三，也是最重要的一点：它不好扩展。这里说的扩展不是指网络做深一点，而是指在“数据规模、算力规模、参数规模”都同步增大时，能不能继续保持稳定的收益曲线。实验已经证明，CNN类骨干达到某个容量之后，继续增加参数得到的提升会明显放缓。

Transformer的统计效率在NLP和视觉分类领域已经被验证过一轮了。它的核心操作是自注意力，每个token可以和任意位置的token交换信息，对全局关系的捕获能力特别强。更关键的是它没有太多卷积归纳偏置，结构统一、模块标准，只需要控制token数量、embedding维度和transformer层数，就能平滑地改变模型规模。DiT要做的，就是把扩散模型中的U-Net骨干整体替换成这种纯Transformer结构，同时保留扩散模型的训练目标。

为了直观对比，我把两者的特性列了一张表：

当然这不等于U-Net一无是处。在中小尺寸、训练数据有限的场景里，U-Net仍然很能打。DiT的意义更多在于：把扩散模型的架构选择问题从“手写网络”变成了“配置化搭建”，这是否是一个更适合扩展的方向，答案在更大计算量上已经比较明显了。

2. DiT的架构设计与关键模块拆解

DiT的整体流程可以拆成四个环节：先把图像切块映射成token序列，然后通过一堆Transformer层做特征变换，在过程中把时间步和类别标签等条件信息注入进去，最后把输出的token序列重新拼回图像尺寸，输出噪声预测。这四个环节听起来简单，但每个环节都有值得琢磨的细节。

2.1 Patchify：图像是怎么变成token的

Transformer处理的是序列，图像是二维格子，怎么把它们变成一维token序列？DiT的做法和Vision Transformer（ViT）一脉相承，叫Patchify。

假设输入是32×32分辨率、3通道的图像，你设定patch大小为8×8，那么每个patch就是一个8×8×3的小立方体。32×32图像被纵向切成4块、横向切成4块，一共得到4×4=16个patch。每个patch通过一个线性层映射为一个embedding向量，于是这16个向量就成了Transformer的16个输入token。token数量就是序列长度，embedding维度就是每个token的向量长度。

从原理层面说，Patchify本质上是在做两件事：一是降维压缩，避免序列长度爆炸；二是把局部特征编码成语义向量。patch大小直接影响性能：patch越小，token数量越多，信息保留越完整，但计算成本也跟着涨。DiT通常使用8或16的patch大小。选8时，32×32输入会得到16个token；选16时，就只剩4个token，虽然快很多，但细节信息丢失得也快。实际做图像生成时，这一步还会考虑操作是否在潜空间进行。

2.2 隐藏在条件注入里的关键设计

纯Transformer只知道输入的token序列，并不知道当前是扩散过程的哪一步，也不知道用户想要生成哪个类别的图片。所以必须额外把这两类信息喂进去。最常见的时间步处理方式是正弦位置编码：把时间步t映射成一个高维向量，再经过MLP变换得到时间嵌入。类别标签则通过嵌入层变成向量。如果训练时使用无分类器引导，类别嵌入还需要在部分样本上置零，模拟“无条件”的情况。

DiT在条件注入上有一次重要探索。社区早期尝试过把条件token拼到输入序列里，让Transformer在自注意力内部自己学会参考；也试过在注意力前面加交叉注意力层，类似语言模型里把文本作为上下文。这些方法都有效，但各自增加了一部分额外计算。DiT最终推荐的方式是adaLN，即自适应层归一化。

2.3 为什么adaLN-zero是主角

传统的层归一化是对特征做减均值、除方差，再通过可学习的缩放参数γ和偏移参数β进行仿射变换。adaLN的思想是：不再用固定的γ和β，而是让条件信息经过一个MLP后动态生成这些参数。每个Transformer块都会接收一组自己的γ和β，也就是说每个块都能根据当前的时间步和类别信息，自主决定应该如何缩放和偏移特征。

adaLN-zero则更进了一步。它把生成γ、β之外，还额外生成一个门控参数α，然后在残差连接的末端加上一层“门控缩放”。最关键的是，这个MLP最后一个线性层在初始化时故意被置为全零输出，那么初始状态下α=0，残差块的输出完全就是恒等映射。这个设计非常巧妙：网络初期从一个稳定的“什么都不做”状态开始，然后逐步学习复杂的变换。这比初始就从随机状态硬学要稳定得多，收敛速度和最终效果都有明显提升。

用代码表达adaLN调制部分大致长这样：

class AdaDNZero(nn.Module): def __init__(self, hidden_size, cond_size): super().__init__() self.net = nn.Sequential( nn.Linear(cond_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size) ) # 最后一个线性层初始化为0 nn.init.zeros_(self.net[-1].weight) nn.init.zeros_(self.net[-1].bias) def forward(self, c): return self.net(c) # 输出按顺序分成6组参数

2.4 四种条件注入方案对比

DiT论文里系统对比过四种方案，我把它们拉开对比一下：

从实验数据看，在同等算力和训练步数下，adaLN-zero的FID指标普遍优于前几种方案。它直接成了实现DiT的首选方式。很多后来出现的生成模型，也都沿用了这种“条件信息通过MLP映射到各层仿射参数”的设计范式。

2.5 解码：把token还原成图像

Transformer输出的是序列化的特征向量，但下游任务需要的是像素级别的噪声预测。最后需要一个解码模块，把它重塑回空间维度。

做法很简单：在最后一个Transformer块之后，先通过层归一化，再过一层线性层，把每个位置的embedding映射成patch_size × patch_size × 输出通道数。这里的输出通道数通常是输入图像通道数的两倍，因为官方实现采用了v-prediction预测方式，每个像素要输出一个预测噪声分量和一个预测图像分量。最后把token序列重新排成网格，得到完整的预测图。

解码部分代码框架：

class FinalLayer(nn.Module): def __init__(self, hidden_size, patch_size, out_channels, cond_size): super().__init__() self.norm_final = nn.LayerNorm(hidden_size) self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(cond_size, 2 * hidden_size) ) def forward(self, x, c): scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1) x = self.norm_final(x) * (1 + scale) + shift x = self.linear(x) return x

这里直接用线性层把特征投影回像素空间，避免了传统解码器里复杂的反卷积堆叠。Transformer的最后一层输出是全局语义特征，线性投影就是一个“字典查找”的过程，把语义映射回局部图案。实际操作中如果觉得画面不够精细，也可以额外加一对反卷积层做上采样，尤其是当patch尺寸偏大时。

3. 核心代码实现：从零搭建一个可运行的DiT

下面进入实战部分。我会用PyTorch从零写一个足够自己训练小规模数据的DiT实现。代码目标是让结构清晰、容易改参数，而不是追求极致性能优化。

3.1 环境与基础配置

需要PyTorch 2.0以上，配合CUDA环境。建议安装timm库用于部分组件。

pip install torch timm einops

代码中主要用到的常量是模型宽度的缩放系数。DiT有S、B、L、XL等配置，它们的区别主要在hidden_size、深度层数和注意力头数。我们以类似DiT-B的配置为例：hidden_size为384，Transformer层数为12，注意力头数为6。

3.2 基础组件：时间嵌入、标签嵌入、PatchEmbed

时间步嵌入沿用正弦位置编码方式，再加一个MLP将它映射到隐藏空间，这个设计已经被无数实验验证过，没什么好改的。

import torch import torch.nn as nn import math # 时间步嵌入 class TimestepEmbedder(nn.Module): def __init__(self, hidden_size, frequency_embedding_size=256): super().__init__() self.mlp = nn.Sequential( nn.Linear(frequency_embedding_size, hidden_size), nn.SiLU(), nn.Linear(hidden_size, hidden_size) ) self.frequency_embedding_size = frequency_embedding_size @staticmethod def timestep_embedding(t, dim, max_period=10000): half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(half, dtype=torch.float32) / half ).to(t.device) args = t[:, None] * freqs[None] return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) def forward(self, t, dtype=torch.float32): t_freq = self.timestep_embedding(t, self.frequency_embedding_size) return self.mlp(t_freq.to(dtype))

类别标签嵌入处理起来会麻烦一点。训练时如果你用classifier-free guidance，需要有一定概率把标签置成“空”类别，也就是让模型在训练时见过无条件的样本。这里我把标签的dropout概率设为0.1。

class LabelEmbedder(nn.Module): def __init__(self, num_classes, hidden_size, drop_prob=0.1): super().__init__() self.embedding = nn.Embedding(num_classes, hidden_size) self.num_classes = num_classes self.drop_prob = drop_prob def forward(self, labels, train=False): if train and self.drop_prob > 0: drop_mask = torch.rand(labels.shape[0], device=labels.device) < self.drop_prob labels = torch.where(drop_mask, torch.full_like(labels, 0, dtype=torch.long), labels) return self.embedding(labels)

PatchEmbed稍微提两句，我把输入通道数设为可配置，默认按灰度图或单通道噪声图来，实际用彩色图时改成3即可。

class PatchEmbed(nn.Module): def __init__(self, input_size=32, patch_size=8, in_channels=3, hidden_size=384): super().__init__() self.patch_size = patch_size self.num_patches = (input_size // patch_size) ** 2 self.proj = nn.Conv2d(in_channels, hidden_size, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x: [B, C, H, W] x = self.proj(x) # [B, hidden, H/P, W/P] x = x.flatten(2).transpose(1, 2) # [B, num_patches, hidden] return x

3.3 注意力机制与DiTBlock

DiT用的注意力是标准的多头自注意力。为了避免画面出现奇怪的伪影，DiT在注意力层里还会做QK归一化，也就是对query和key先做归一化再算点积，这样可以稳定训练。我在这里加上这一细节。

class Attention(nn.Module): def __init__(self, hidden_size, num_heads, qk_norm=True): super().__init__() self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=True) self.proj = nn.Linear(hidden_size, hidden_size, bias=True) if qk_norm: self.q_norm = nn.LayerNorm(self.head_dim, elementwise_affine=True) self.k_norm = nn.LayerNorm(self.head_dim, elementwise_affine=True) else: self.q_norm = nn.Identity() self.k_norm = nn.Identity() def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] q = self.q_norm(q) k = self.k_norm(k) attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1, 2).reshape(B, N, C) return self.proj(x)

核心的DiTBlock结构是：先做adaLN-zero调制，然后过注意力块，再过MLP，最后门控缩放叠加到残差上。

class DiTBlock(nn.Module): def __init__(self, hidden_size, num_heads, cond_size, mlp_ratio=4.0): super().__init__() self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.attn = Attention(hidden_size, num_heads) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(hidden_size, int(hidden_size * mlp_ratio)), nn.GELU(approximate='tanh'), nn.Linear(int(hidden_size * mlp_ratio), hidden_size) ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(cond_size, 6 * hidden_size, bias=True) ) # 零初始化最后一个线形层 nn.init.zeros_(self.adaLN_modulation[-1].weight) nn.init.zeros_(self.adaLN_modulation[-1].bias) def forward(self, x, c): scale1, shift1, gate1, scale2, shift2, gate2 = self.adaLN_modulation(c).chunk(6, dim=-1) h = self.norm1(x) * (1 + scale1) + shift1 h = self.attn(h) x = x + gate1.unsqueeze(1) * h h = self.norm2(x) * (1 + scale2) + shift2 h = self.mlp(h) x = x + gate2.unsqueeze(1) * h return x

我特别想把第六个参数gate2单独提出来说：它专门控制MLP分支的权重，在初始化全零的情况下，每个块一开始都相当于一个“直通”，这种设计让深层网络训练时前期风险降到了最低。

3.4 整体模型组合

有了上面的组件，组装DiT模型就很简单了。forward里做的事是：patch化、加位置编码、过一系列DiTBlock，最后过FinalLayer把token映射回像素。

class DiT(nn.Module): def __init__(self, input_size=32, patch_size=8, in_channels=3, hidden_size=384, depth=12, num_heads=6, num_classes=10, learn_sigma=True): super().__init__() self.in_channels = in_channels self.out_channels = in_channels * 2 if learn_sigma else in_channels self.patch_size = patch_size self.num_heads = num_heads self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size) self.t_embedder = TimestepEmbedder(hidden_size) self.y_embedder = LabelEmbedder(num_classes, hidden_size) num_patches = self.x_embedder.num_patches self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size)) nn.init.trunc_normal_(self.pos_embed, std=0.02) self.blocks = nn.ModuleList([ DiTBlock(hidden_size, num_heads, hidden_size) for _ in range(depth) ]) self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels, hidden_size) def forward(self, x, t, y): x = self.x_embedder(x) x = x + self.pos_embed t = self.t_embedder(t) y = self.y_embedder(y) c = t + y for block in self.blocks: x = block(x, c) x = self.final_layer(x, c) x = x.reshape(x.shape[0], self.out_channels, -1, self.patch_size, self.patch_size) # 从token序列重排回图像网格 B = x.shape[0] C_out = self.out_channels P = self.patch_size H = W = int(x.shape[2] ** 0.5) x = x.reshape(B, C_out, H, P, W, P) x = x.permute(0, 1, 2, 4, 3, 5).reshape(B, C_out, H * P, W * P) return x

这里值得注意的是我把位置编码作为可学习参数初始化。相比固定正弦位置编码，可学习编码在小尺寸输入上往往表现更好，因为它可以直接被梯度优化到最适合当前patch布局的位置分布。

3.5 训练与采样的核心逻辑

训练循环本身并不复杂，关键是构造带噪图像。给定一批干净图片x_0和随机时间步t，先生成噪声ε，再通过一个预计算好的噪声系数把二者混合，得到x_t。模型的目标是尽可能准确地预测出这个ε。

def training_step(model, x_0, y, t, noise, betas, alphas_cumprod): sqrt_alpha = torch.sqrt(alphas_cumprod[t]).view(-1, 1, 1, 1) sqrt_one_minus_alpha = torch.sqrt(1 - alphas_cumprod[t]).view(-1, 1, 1, 1) x_t = sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise pred = model(x_t, t, y) # 如果模型输出两个通道，取前一半作为噪声预测 if pred.shape[1] == 2 * model.in_channels: pred = pred[:, :model.in_channels] loss = torch.nn.functional.mse_loss(pred, noise) return loss

采样则使用DDPM的逻辑，先从一个纯噪声开始，然后从t=T逐步往t=0走，每一步都让模型预测噪声，把当前样本稍微去噪一点，再补一点随机扰动。

@torch.no_grad() def sample(model, shape, num_steps, alphas_cumprod, betas, labels=None, guidance_scale=0.0, device='cuda'): x = torch.randn(shape, device=device) for i in range(num_steps - 1, -1, -1): t = torch.full((shape[0],), i, device=device, dtype=torch.long) alpha = alphas_cumprod[i] alpha_next = alphas_cumprod[i - 1] if i > 0 else torch.tensor(1.0) pred_noise = model(x, t, labels) # 无条件 + 条件融合，classifier-free guidance if guidance_scale > 0: null_labels = torch.full_like(labels, 0) pred_uncond = model(x, t, null_labels) pred_noise = pred_uncond + guidance_scale * (pred_noise - pred_uncond) sqrt_alpha = alpha.sqrt() sqrt_one_minus_alpha = (1 - alpha).sqrt() x0_pred = (x - sqrt_one_minus_alpha * pred_noise) / sqrt_alpha # 加入随机采样噪声 noise = torch.randn_like(x) if i > 0 else torch.zeros_like(x) x = alpha_next.sqrt() * x0_pred + (1 - alpha_next).sqrt() * noise return x

看到这里你可能有一个感觉：整体代码量不大，模块很干净。这正是DiT的吸引力——骨干网络的复杂度最高的是注意力，但注意力本身就是标准层，其他的都是可插拔的调制组件。只要你把PatchEmbed和FinalLayer的尺寸算清楚，在其他图像生成任务里复用这套代码的成本非常低。

4. 训练技巧、常见问题与排查实录

这篇代码我实际跑过的次数不算少，训练和采样阶段踩过的坑也值得单独整理出来。这些问题网上零星有一些讨论，但系统性整理成速查表的很少。

4.1 模型一直在学，loss就是不降

如果你发现自己训的DiT loss很久都不降，先别急着改模型结构，按顺序排查这几件事：

第一，检查时间步t是否真的均匀采样。扩散模型的训练目标是在所有时间步上等概率地随机采样，如果你用了一个偏向低噪声的采样分布，模型会完全忽略高噪声区间的学习，loss曲线会一直拖着一条“尾巴”不往下走。第二，查看你的模型输入输出通道是否匹配。DiT如果采用v-prediction，输出通道数是输入的两倍，很多人在封装模型时忘了这个细节，导致loss计算时张量一直在错误地广播。第三，学习率不是越小越好。DiT一般用AdamW优化器，基础学习率设置在1e-4到2e-4之间，配合线性warmup和余弦退火。如果设置成1e-5，模型前期根本学不动。

还有一个容易被忽略的点，就是位置编码的标准差。DiT把位置编码初始化为均值为0、标准差0.02的正态分布，这个值来自经验。如果初始化标准差太大，位置信息一开始就是紊乱的，注意力很难看到合理的局部空间结构。我的建议是固定用trunc_normal，并且把标准差控制在0.02这个量级。

4.2 生成图像颜色发灰、边缘模糊怎么办

颜色问题首先想到classifier-free guidance的引导强度。引导强度太小，生成结果中类别特征不鲜明，画面偏灰偏平滑；引导强度太大，又会出现过饱和和像素伪影。实践下来，在类别引导任务上，2到4是一个比较舒服的范围，但具体还得看你的数据集和类别embedding的落差。

模糊问题则要关注采样步数。DDPM采样如果只用10步，基本就是自然模糊，建议先用50步以上来排除采样步数的影响。如果你的训练集本身是小尺寸图像，生成大图时很容易出现模糊，这是因为patch大小没有随分辨率调整。原本用32×32训练的模型，如果直接拿来生成64×64，token数变了，位置编码也可能不匹配。正确的做法是保持patch size一致，让序列长度自然翻倍，但这样需要重新训练或者做位置编码插值。

4.3 显存爆掉的常见原因

DiT虽然结构统一，但显存压力一点也不小。第一需要注意序列长度：patch越小token越多，随着图片分辨率提升，attention占用的显存是平方级增长的。32×32的图像，patch=8时有16个token；放到256×256，patch=8时变成1024个token，这比多数NLP任务还要紧张。

常用的优化手段有：

• 混合精度训练：把模型计算切成FP16或BF16，DiT对精度不敏感，实测BF16效果很好。
• 梯度检查点：attention和MLP的前向输出不保留，反向时重新计算，能省不少显存，代价是训练时间变长。
• 减少注意力头数：只要通道维度被头数整除，就能正常跑，如果显存吃紧，可以把头数从8降到4，效果损失相对可控。

还有一个经验值得单独说：在采样阶段使用fp32推理。混合精度训练没问题，但推理时如果直接沿用fp16，attention的softmax在极端分数下容易溢出，产生画面上的黑斑或白斑。我自己的习惯是训练用BF16，采样单独走全精度。

4.4 代码常见报错速查表

4.5 一个从零开始的训练经验清单

我把自己实际操作中最顺手的配置列一下：使用ImageNet的类别条件作为示例任务，输入分辨率32×32，patch=8，hidden_size=384，depth=12，head=6，batch size设为128。训练使用AdamW，初始学习率1e-4，weight decay设为0.3。整个训练过程大概在几千步后就能看到比较清晰的类别分化。

为什么weight decay要设这么大？DiT是纯Transformer，参数众多且高度依赖归一化。过小的weight decay会让attention的残差输出逐渐累积，导致深层网络训练不稳定。0.3到0.5这个区间是我测试下来最稳的范围，尤其适合训练步数较长的扩散模型。

另一个容易影响情绪的点是训练曲线不会像分类任务那样快速收敛。扩散模型要分阶段：前期loss缓慢下降，中期loss稳步下滑但指标提升很快。如果一定要用训练曲线来守护判断，我建议不要只看loss，而是每隔一定步数真正采样一批图片出来看，眼见的画面比曲线数据可靠得多。

5. 最后想分享的一点工程心得

写完这整套流程，最深的体会可能和大多数人预想的不太一样。DiT最费时间的地方，并不在理解Transformer原理，也不在写模型主体代码，而在于把条件注入、patch尺寸、通道数、采样参数这些容易互相牵连的细节对齐。我见过不只一个人把模型搭起来、训练loss正常下降，但最后采样生成的图片怎么看都像“打码”——最后发现是patch重排回图像时维度的顺序错了。所以如果你也是照着这篇内容自己写一遍，强烈建议在训练前先用一张真实图片过一遍前向流程，把输出尺寸和预期值打印出来逐一比对。

扩散模型这几年迭代速度很快，DiT本身也带了很多后续工作。我自己的习惯是，学一个新模型时必须动手写一遍核心代码，写不完理解始终是虚的。这篇文章给出的实现足够在小规模任务上跑通全过程，如果你接着想往上加一些新想法，比如更长的序列注意力、多模态条件注入、更高效的采样器，都可以在这套结构上无缝地换层。希望这篇东西能帮你省下一些翻论文和调bug的时间。