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

发布时间:2026/7/24 17:06:48

Stable Diffusion 源码学习:VAE 编码器、卷积、残差块与重参数化采样 Stable Diffusion 源码学习VAE 编码器、卷积、残差块与重参数化采样前言Stable Diffusion 并不是直接在原始图片的像素空间中进行扩散而是先通过 VAE 编码器将图片压缩成潜变量latent然后在潜空间中完成加噪和去噪。本文结合 PyTorch Stable Diffusion 项目中的 VAE 代码整理以下知识点nn.Module和nn.Sequentialnn.Conv2d参数及可学习的卷积核VAE_ResidualBlock残差块编码器中的下采样与非对称 paddinggetattr()和卷积步长元组什么是 VAE均值、方差与高维潜变量重参数化采样最终编码得到的z是什么一、nn.Module与nn.Sequential1.nn.Modulenn.Module是 PyTorch 中所有神经网络模块的基础类。自定义网络时通常需要在__init__()中定义网络层并在forward()中定义数据如何流动importtorch.nnasnnclassMyModel(nn.Module):def__init__(self):super().__init__()self.linear1nn.Linear(10,20)self.relunn.ReLU()self.linear2nn.Linear(20,2)defforward(self,x):xself.linear1(x)xself.relu(x)xself.linear2(x)returnxnn.Module可以实现顺序结构、残差连接、条件分支、多输入和多输出等复杂计算流程。2.nn.Sequentialnn.Sequential本身也是一个nn.Module但它专门用来组织按顺序执行的网络层modelnn.Sequential(nn.Linear(10,20),nn.ReLU(),nn.Linear(20,2))执行outputmodel(x)相当于xmodel[0](x)xmodel[1](x)outputmodel[2](x)数据流为输入 → Linear → ReLU → Linear → 输出对比项nn.Sequential自定义nn.Module是否需要编写forward()通常不需要需要数据流固定顺序执行可以自由控制条件分支不方便支持残差连接不方便支持多输入、多输出不方便支持适用场景简单顺序网络复杂网络可以把nn.Module理解成搭建神经网络的基础框架而nn.Sequential是快速搭建顺序网络的容器。二、理解nn.Conv2d代码nn.Conv2d(3,128,kernel_size3,padding1)完整写法是nn.Conv2d(in_channels3,out_channels128,kernel_size3,stride1,padding1)参数当前值含义in_channels3输入通道数RGB 图片有 3 个通道out_channels128输出 128 个特征通道kernel_size3卷积核尺寸为3×3stride1卷积核每次移动 1 个像素padding1在输入四周补零假设输入张量形状为(B, 3, H, W)输出形状为(B, 128, H, W)其中B是批次大小H是图片高度W是图片宽度。1. 卷积输出尺寸卷积输出尺寸公式为H_out floor((H_in 2P - K) / S) 1其中PpaddingKkernel sizeSstride代入kernel_size3、padding1、stride1H_out floor((H 2 - 3) / 1) 1 H因此图片的高度和宽度不变只将通道数从 3 变成 128。2. 卷积核参数是否需要学习需要。卷积层包含可学习的权重和偏置convnn.Conv2d(3,128,kernel_size3,padding1)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()outputmodel(x)losscriterion(output,target)loss.backward()optimizer.step()kernel_size3只规定卷积核尺寸为3×3卷积核内部的数值是训练得到的。三、VAE_ResidualBlock是什么VAE_ResidualBlock不是 PyTorch 内置模块而是项目自定义的 VAE 残差块。classVAE_ResidualBlock(nn.Module):def__init__(self,in_channels,out_channels):super().__init__()self.groupnorm_1nn.GroupNorm(32,in_channels)self.conv_1nn.Conv2d(in_channels,out_channels,kernel_size3,padding1)self.groupnorm_2nn.GroupNorm(32,out_channels)self.conv_2nn.Conv2d(out_channels,out_channels,kernel_size3,padding1)ifin_channelsout_channels:self.residual_layernn.Identity()else:self.residual_layernn.Conv2d(in_channels,out_channels,kernel_size1)defforward(self,x):residuex xself.groupnorm_1(x)xF.silu(x)xself.conv_1(x)xself.groupnorm_2(x)xF.silu(x)xself.conv_2(x)returnxself.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):xF.pad(x,(0,1,0,1))xmodule(x)因为编码器继承了nn.Sequential所以self中保存了编码器的所有网络层classVAE_Encoder(nn.Sequential):...循环相当于xself[0](x)xself[1](x)xself[2](x)# ...1.getattr()的作用getattr(module,stride,None)格式为getattr(对象,属性名,默认值)如果module有stride属性就返回module.stride否则返回None。这样可以避免直接访问不存在的属性而报错。2. 为什么stride是元组二维卷积需要分别描述高度和宽度方向的步长stride (高度方向步长, 宽度方向步长)即使传入的是整数nn.Conv2d(...,stride2)PyTorch 内部也会将其转换为(2,2)因此getattr(module,stride,None)(2,2)是在判断当前模块是否为高度和宽度方向步长都为 2 的层。该表达式整体返回True或False。五、为什么下采样前要补零xF.pad(x,(0,1,0,1))对于二维图像四个数字依次表示(左边, 右边, 上边, 下边)因此该操作会在输入右边补 1 列零在下边补 1 行零。假设卷积为nn.Conv2d(in_channels,out_channels,kernel_size3,stride2,padding0)输入尺寸为 512不补零时H_out floor((512 - 3) / 2) 1 255补零后尺寸由 512 变成 513H_out floor((513 - 3) / 2) 1 256整个过程为512×512 ↓ 右边和下边补零 513×513 ↓ kernel_size3stride2 256×256直接设置padding1也能得到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 AutoencoderVAE不会让编码器直接输出唯一确定的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_variancetorch.chunk(x,2,dim1)编码器输出(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 中可以使用epsilontorch.randn(shape)生成标准高斯噪声。十一、重参数化采样VAE 使用以下公式获得潜变量z μ σε其中ε ~ N(0, I)数学符号代码变量含义μmean编码器预测的均值log(σ²)log_variance编码器预测的对数方差σstdev标准差εnoise标准高斯噪声zz最终潜变量对应代码mean,log_variancetorch.chunk(x,2,dim1)variancelog_variance.exp()stdevvariance.sqrt()zmeanstdev*noise1. 噪声来自哪个分布ε来自事先设定的标准高斯分布ε ~ 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。十六、图生图与文生图中的 VAE1. 图生图图生图具有输入图片因此需要先编码输入图片 ↓ VAE Encoder 潜变量z ↓ 加入扩散噪声 带噪潜变量 ↓ UNet逐步去噪 新潜变量 ↓ VAE Decoder 输出图片2. 文生图文生图没有输入图片所以通常不需要在开始时调用 VAE Encoder随机潜空间噪声 ↓ UNet根据文本逐步去噪 生成的潜变量 ↓ VAE Decoder 最终图片十七、VAE 采样噪声与扩散噪声的区别VAE 中的噪声和扩散过程中的噪声通常都来自标准高斯分布但用途不同。VAE 中的噪声用于从图片对应的潜变量分布中获得zz μ σε扩散过程中的噪声用于逐渐污染潜变量然后让 UNet 学习如何去除噪声。噪声用途VAE 采样噪声从 q(z扩散噪声对潜变量加噪并由 UNet 去噪二者虽然可能都服从标准高斯分布但属于两个不同阶段。总结整个 VAE 编码过程可以概括为x → 卷积编码器 → μ(x), log(σ²(x)) → ε ~ N(0,I) → z μ σε各部分职责如下卷积和残差块 从图片中提取特征并降低空间分辨率 最后一层卷积 输出μ和log(σ²) 标准高斯采样 生成ε 重参数化 通过zμσε得到潜变量 最终的z 图片的压缩特征表示同时具有受控随机性最核心的理解是卷积编码器决定潜变量应该位于什么位置以及允许多大的变化标准高斯噪声决定本次具体取到哪个位置最终得到的z是用于图片重建和潜空间扩散的压缩特征。

相关新闻