Stable Diffusion 源码学习：VAE 编码器、卷积、残差块与重参数化采样

Stable Diffusion 源码学习：VAE 编码器、卷积、残差块与重参数化采样

前言

Stable Diffusion 并不是直接在原始图片的像素空间中进行扩散，而是先通过 VAE 编码器将图片压缩成潜变量（latent），然后在潜空间中完成加噪和去噪。

本文结合 PyTorch Stable Diffusion 项目中的 VAE 代码，整理以下知识点：

1. nn.Module和nn.Sequential
2. nn.Conv2d参数及可学习的卷积核
3. VAE_ResidualBlock残差块
4. 编码器中的下采样与非对称 padding
5. getattr()和卷积步长元组
6. 什么是 VAE
7. 均值、方差与高维潜变量
8. 重参数化采样
9. 最终编码得到的z是什么

一、nn.Module与nn.Sequential

1.nn.Module

nn.Module是 PyTorch 中所有神经网络模块的基础类。

自定义网络时，通常需要在__init__()中定义网络层，并在forward()中定义数据如何流动：

importtorch.nnasnnclassMyModel(nn.Module):def__init__(self):super().__init__()self.linear1=nn.Linear(10,20)self.relu=nn.ReLU()self.linear2=nn.Linear(20,2)defforward(self,x):x=self.linear1(x)x=self.relu(x)x=self.linear2(x)returnx

nn.Module可以实现顺序结构、残差连接、条件分支、多输入和多输出等复杂计算流程。

2.nn.Sequential

nn.Sequential本身也是一个nn.Module，但它专门用来组织按顺序执行的网络层：

model=nn.Sequential(nn.Linear(10,20),nn.ReLU(),nn.Linear(20,2))

执行：

output=model(x)

相当于：

x=model[0](x)x=model[1](x)output=model[2](x)

数据流为：

输入 → Linear → ReLU → Linear → 输出

可以把nn.Module理解成搭建神经网络的基础框架，而nn.Sequential是快速搭建顺序网络的容器。

二、理解nn.Conv2d

代码：

nn.Conv2d(3,128,kernel_size=3,padding=1)

完整写法是：

nn.Conv2d(in_channels=3,out_channels=128,kernel_size=3,stride=1,padding=1)

假设输入张量形状为：

(B, 3, H, W)

输出形状为：

(B, 128, H, W)

其中B是批次大小，H是图片高度，W是图片宽度。

1. 卷积输出尺寸

卷积输出尺寸公式为：

H_out = floor((H_in + 2P - K) / S) + 1

其中：

• P：padding
• K：kernel size
• S：stride

代入kernel_size=3、padding=1、stride=1：

H_out = floor((H + 2 - 3) / 1) + 1 = H

因此图片的高度和宽度不变，只将通道数从 3 变成 128。

2. 卷积核参数是否需要学习？

需要。卷积层包含可学习的权重和偏置：

conv=nn.Conv2d(3,128,kernel_size=3,padding=1)print(conv.weight.shape)# torch.Size([128, 3, 3, 3])print(conv.bias.shape)# torch.Size([128])

权重形状表示：

[输出通道数, 输入通道数, 卷积核高度, 卷积核宽度]

权重数量为：

128 × 3 × 3 × 3 = 3456

加上 128 个偏置，总参数量为：

3456 + 128 = 3584

训练过程中，反向传播计算这些参数的梯度，优化器再更新它们：

optimizer.zero_grad()output=model(x)loss=criterion(output,target)loss.backward()optimizer.step()

kernel_size=3只规定卷积核尺寸为3×3，卷积核内部的数值是训练得到的。

三、VAE_ResidualBlock是什么？

VAE_ResidualBlock不是 PyTorch 内置模块，而是项目自定义的 VAE 残差块。

classVAE_ResidualBlock(nn.Module):def__init__(self,in_channels,out_channels):super().__init__()self.groupnorm_1=nn.GroupNorm(32,in_channels)self.conv_1=nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1)self.groupnorm_2=nn.GroupNorm(32,out_channels)self.conv_2=nn.Conv2d(out_channels,out_channels,kernel_size=3,padding=1)ifin_channels==out_channels:self.residual_layer=nn.Identity()else:self.residual_layer=nn.Conv2d(in_channels,out_channels,kernel_size=1)defforward(self,x):residue=x x=self.groupnorm_1(x)x=F.silu(x)x=self.conv_1(x)x=self.groupnorm_2(x)x=F.silu(x)x=self.conv_2(x)returnx+self.residual_layer(residue)

主分支执行：

x → GroupNorm → SiLU → 3×3卷积 → GroupNorm → SiLU → 3×3卷积

最后与残差分支相加：

y = F(x) + R(x)

当输入、输出通道数相同时，例如：

VAE_ResidualBlock(128,128)

使用nn.Identity()保留原输入：

y = F(x) + x

当通道数不同时，例如：

VAE_ResidualBlock(128,256)

使用1×1卷积把残差分支从 128 个通道变成 256 个通道：

y = F(x) + Conv_1×1(x)

残差连接可以保留原始输入特征、改善梯度传播，并让深层网络更容易训练。

四、VAE 编码器中的循环

formoduleinself:ifgetattr(module,"stride",None)==(2,2):x=F.pad(x,(0,1,0,1))x=module(x)

因为编码器继承了nn.Sequential，所以self中保存了编码器的所有网络层：

classVAE_Encoder(nn.Sequential):...

循环相当于：

x=self[0](x)x=self[1](x)x=self[2](x)# ...

1.getattr()的作用

getattr(module,"stride",None)

格式为：

getattr(对象,"属性名",默认值)

如果module有stride属性，就返回module.stride；否则返回None。这样可以避免直接访问不存在的属性而报错。

2. 为什么stride是元组？

二维卷积需要分别描述高度和宽度方向的步长：

stride = (高度方向步长, 宽度方向步长)

即使传入的是整数：

nn.Conv2d(...,stride=2)

PyTorch 内部也会将其转换为：

(2,2)

因此：

getattr(module,"stride",None)==(2,2)

是在判断当前模块是否为高度和宽度方向步长都为 2 的层。该表达式整体返回True或False。

五、为什么下采样前要补零？

x=F.pad(x,(0,1,0,1))

对于二维图像，四个数字依次表示：

(左边, 右边, 上边, 下边)

因此该操作会在输入右边补 1 列零，在下边补 1 行零。

假设卷积为：

nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=2,padding=0)

输入尺寸为 512，不补零时：

H_out = floor((512 - 3) / 2) + 1 = 255

补零后尺寸由 512 变成 513：

H_out = floor((513 - 3) / 2) + 1 = 256

整个过程为：

512×512 ↓ 右边和下边补零 513×513 ↓ kernel_size=3，stride=2 256×256

直接设置padding=1也能得到256×256，但那是对称 padding。项目中采用只在右侧和下侧补零的非对称 padding，是为了与原始 Stable Diffusion VAE 的特征对齐方式和预训练权重保持一致。

六、什么是自编码器？

自编码器由编码器和解码器组成：

图片x → Encoder → 潜变量z → Decoder → 重建图片x'

编码器负责压缩图片，解码器负责根据压缩特征重建图片。例如：

输入图片：(B, 3, 512, 512) ↓ Encoder 潜变量： (B, 4, 64, 64) ↓ Decoder 重建图片：(B, 3, 512, 512)

普通自编码器通常直接输出一个确定的潜变量：

z = Encoder(x)

七、什么是变分自编码器？

变分自编码器（Variational Autoencoder，VAE）不会让编码器直接输出唯一确定的z，而是输出潜变量概率分布的参数。

通常假设潜变量服从高斯分布：

q(z|x) = N(μ, σ²)

编码器根据输入图片x计算：

• 均值μ
• 对数方差log(σ²)

然后通过采样得到潜变量z：

图片x ↓ 卷积编码器 μ 和 log(σ²) ↓ 重参数化采样 潜变量z ↓ Decoder 重建图片

八、卷积过程和采样过程有什么联系？

卷积和采样是前后两个阶段：

卷积负责分析并压缩图片，产生分布参数；采样根据这些参数产生最终潜变量。

对于一张512×512的图片，形状变化大致为：

(B, 3, 512, 512) ↓ 卷积 (B, 128, 512, 512) ↓ 下采样 (B, 128, 256, 256) ↓ 下采样 (B, 256, 128, 128) ↓ 下采样 (B, 512, 64, 64) ↓ 最后一层卷积 (B, 8, 64, 64)

最后的 8 个通道被分成两组：

mean,log_variance=torch.chunk(x,2,dim=1)

编码器输出：(B, 8, 64, 64) │ ├── 前4个通道 → mean = μ │ (B, 4, 64, 64) │ └── 后4个通道 → log_variance = log(σ²) (B, 4, 64, 64)

不是mean的前 4 个通道，而是编码器输出张量的前 4 个通道共同组成mean。

九、均值和方差也是高维张量

μ和log(σ²)都不是单个数字，而是高维张量：

μ： (B, 4, 64, 64) log(σ²)： (B, 4, 64, 64)

对于一个样本，共有：

4 × 64 × 64 = 16384

个潜变量维度。编码器会为每个维度分别预测一个均值和方差：

μ[b,c,h,w]

σ²[b,c,h,w]

通常假设给定输入图片后，各潜变量维度条件独立，因此使用对角高斯分布：

q(z|x) = N(μ, diag(σ²))

模型只预测每个维度自己的方差，不需要预测巨大的完整协方差矩阵。

十、什么是采样？

采样可以理解为按照某个概率分布，随机抽取一个数或张量。

ε ~ N(0, I)

表示从预先规定的标准高斯分布中采样随机噪声ε。

在一维情况下：

ε ~ N(0, 1)

多次采样可能得到0.2、-1.1、0.7、-0.3等数值，它们通常集中在 0 附近。

在 PyTorch 中可以使用：

epsilon=torch.randn(shape)

生成标准高斯噪声。

十一、重参数化采样

VAE 使用以下公式获得潜变量：

z = μ + σε

其中：

ε ~ N(0, I)

对应代码：

mean,log_variance=torch.chunk(x,2,dim=1)variance=log_variance.exp()stdev=variance.sqrt()z=mean+stdev*noise

1. 噪声来自哪个分布？

ε来自事先设定的标准高斯分布：

ε ~ N(0, I)

它不是从编码器产生的q(z|x)中采样的。

编码器计算出μ和σ后，将标准高斯噪声变换为：

z = μ + σε

变换后的z等价于从下面的分布中采样：

z ~ q(z|x) = N(μ, σ²)

也就是说：

标准高斯分布N(0,I) ↓ 采样 ε ↓ 乘以σ，再加上μ z ↓ z服从N(μ,σ²)

2. 数值示例

假设编码器得到：

μ = 10, σ = 2

从标准高斯分布采样得到：

ε = 0.5

那么：

z = 10 + 2 × 0.5 = 11

这里0.5是标准高斯噪声，11是最终的潜变量样本。由于σ² = 4，对应分布为N(10, 4)。

3. 高维情况下的计算

实际代码中的变量都是形状相同的张量：

mean： (B, 4, 64, 64) stdev： (B, 4, 64, 64) noise： (B, 4, 64, 64) z： (B, 4, 64, 64)

计算逐元素进行：

z[b,c,h,w] = μ[b,c,h,w] + σ[b,c,h,w]ε[b,c,h,w]

十二、为什么需要重参数化？

如果直接把采样写成：

z ~ N(μ, σ²)

随机采样操作会让反向传播变得困难。VAE 把随机性单独放在与模型参数无关的ε中：

z = μ + σε

此时z对μ和σ可微：

∂z/∂μ = 1

∂z/∂σ = ε

梯度传播过程为：

损失 ↓ z ↓ μ和σ ↓ 生成μ和σ的卷积层 ↓ 更新卷积核权重和偏置

这就是 VAE 中的重参数化技巧。

十三、μ和σ是可学习参数吗？

严格来说，它们不是模型中直接保存的固定参数。真正由优化器更新的是编码器中的卷积核权重、偏置和其他网络层参数。

μ和σ是编码器针对当前输入图片计算出来的结果：

图片A → μA、σA 图片B → μB、σB

不同图片会得到不同的μ和σ，但使用的是同一套编码器权重。

虽然μ和σ不是固定的模型参数，但它们参与计算图，梯度会经过它们传回编码器。

十四、最后编码得到的z是什么？

最终的z是输入图片在 VAE 潜空间中的压缩特征表示，同时带有采样引入的一定随机性。

z = μ(x) + σ(x)ε

其中：

• μ(x)：图片潜表示的中心，包含主要图片特征；
• σ(x)：编码器预测的不确定程度；
• ε：标准高斯随机噪声；
• z：本次采样得到的图片潜表示。

如果σ很小：

z ≈ μ

同一张图片多次编码得到的z会非常接近。如果σ较大，随机变化会更加明显。

同一个输入x ├── ε₁ → z₁ = μ(x) + σ(x)ε₁ └── ε₂ → z₂ = μ(x) + σ(x)ε₂

一般来说z₁ ≠ z₂，但它们都围绕同一个μ(x)，因此都应保留原图片的主要信息。

潜变量可能包含：

• 图片整体结构
• 物体形状
• 空间布局
• 颜色和纹理
• 边缘与局部特征
• Decoder 重建图片所需的信息

不能简单认为某个通道只表示颜色、另一个通道只表示轮廓。它采用的是分布式特征表示，单个数值通常没有明确的人类语义。

十五、为什么 Stable Diffusion 使用 VAE？

对于512×512的图片：

原图片：(B, 3, 512, 512) 潜变量：(B, 4, 64, 64)

原图元素数量为：

3 × 512 × 512 = 786432

潜变量元素数量为：

4 × 64 × 64 = 16384

两者相差：

786432 / 16384 = 48

因此潜变量的数据量大约是原始图片的1/48。Stable Diffusion 在更小的潜空间中运行扩散模型，可以显著减少计算量和显存占用。

图片 ↓ VAE Encoder 潜变量z ↓ 扩散模型加噪和去噪 新的潜变量z' ↓ VAE Decoder 输出图片

这种方法称为潜空间扩散（Latent Diffusion）。

十六、图生图与文生图中的 VAE

1. 图生图

图生图具有输入图片，因此需要先编码：

输入图片 ↓ VAE Encoder 潜变量z ↓ 加入扩散噪声 带噪潜变量 ↓ UNet逐步去噪 新潜变量 ↓ VAE Decoder 输出图片

2. 文生图

文生图没有输入图片，所以通常不需要在开始时调用 VAE Encoder：

随机潜空间噪声 ↓ UNet根据文本逐步去噪 生成的潜变量 ↓ VAE Decoder 最终图片

十七、VAE 采样噪声与扩散噪声的区别

VAE 中的噪声和扩散过程中的噪声通常都来自标准高斯分布，但用途不同。

VAE 中的噪声用于从图片对应的潜变量分布中获得z：

z = μ + σε

扩散过程中的噪声用于逐渐污染潜变量，然后让 UNet 学习如何去除噪声。

二者虽然可能都服从标准高斯分布，但属于两个不同阶段。

总结

整个 VAE 编码过程可以概括为：

x → 卷积编码器 → μ(x), log(σ²(x)) → ε ~ N(0,I) → z = μ + σε

各部分职责如下：

卷积和残差块： 从图片中提取特征并降低空间分辨率 最后一层卷积： 输出μ和log(σ²) 标准高斯采样： 生成ε 重参数化： 通过z=μ+σε得到潜变量 最终的z： 图片的压缩特征表示，同时具有受控随机性

最核心的理解是：

卷积编码器决定潜变量应该位于什么位置，以及允许多大的变化；标准高斯噪声决定本次具体取到哪个位置；最终得到的z是用于图片重建和潜空间扩散的压缩特征。