VAE：原理+代码全解析

变分自编码器Variational AutoencoderVAE是深度学习中经典的生成模型之一它结合了自编码器的结构和变分推断的思想既能完成数据压缩又能实现数据生成。本文将从原理到代码一步步拆解VAE的核心逻辑。一、VAE的核心思想VAE的本质是通过学习数据的潜在分布实现从低维隐空间到高维数据空间的映射。和传统自编码器不同VAE不是直接学习输入到隐向量的确定性映射而是学习隐向量的概率分布这也是它能生成新数据的关键。1. 传统自编码器的局限传统自编码器由编码器和解码器组成编码器将输入数据压缩成固定维度的隐向量解码器再将隐向量还原为输入数据。但这种结构的隐空间是离散且无规律的无法通过采样隐向量生成新数据——比如在两个隐向量之间插值可能得到无意义的结果。2. VAE的改进引入概率分布VAE对编码器做了修改不再输出固定的隐向量而是输出隐向量的均值μ和方差σ²为了计算方便通常输出logσ²避免方差为负。然后从这个正态分布N(μ, σ²)中采样得到隐向量z再输入解码器还原数据。这个过程可以用两个核心步骤概括编码过程输入x → 编码器输出μ和logσ² → 采样得到z ~ N(μ, σ²)解码过程z → 解码器输出重构数据x̂3. VAE的损失函数VAE的损失由两部分组成重构损失和KL散度损失。1重构损失衡量解码器输出的重构数据x̂和原始输入x的差异通常用交叉熵损失针对图像等离散数据或均方误差针对连续数据Lrecon−Ez∼q(z∣x)[log⁡p(x∣z)]L_{recon} -\mathbb{E}_{z \sim q(z|x)}[\log p(x|z)]Lrecon​−Ez∼q(z∣x)​[logp(x∣z)]简单来说就是让重构数据尽可能接近原始数据。2KL散度损失KL散度用于衡量编码器输出的分布q(z|x)和预设的先验分布p(z)通常设为标准正态分布N(0,1)之间的差异LKLDKL(q(z∣x)∣∣p(z))12∑i1d(μi2σi2−log⁡σi2−1)L_{KL} D_{KL}(q(z|x) || p(z)) \frac{1}{2}\sum_{i1}^d (\mu_i^2 \sigma_i^2 - \log\sigma_i^2 - 1)LKL​DKL​(q(z∣x)∣∣p(z))21​i1∑d​(μi2​σi2​−logσi2​−1)这部分损失的作用是约束隐空间的分布尽可能接近标准正态分布保证隐空间的连续性和规律性这样在隐空间中采样就能生成有意义的数据。最终VAE的总损失为LLreconLKLL L_{recon} L_{KL}LLrecon​LKL​二、重参数化技巧这里有个关键问题如果直接从N(μ, σ²)中采样z反向传播时梯度无法通过采样操作传递因为采样是随机过程不可导。为了解决这个问题VAE引入了重参数化技巧将采样过程改写为zμσ⊙ϵ,ϵ∼N(0,1)z \mu \sigma \odot \epsilon, \quad \epsilon \sim N(0,1)zμσ⊙ϵ,ϵ∼N(0,1)其中⊙表示元素-wise乘法。这样一来采样的随机性转移到了ε上而μ和σ是编码器的输出梯度可以通过μ和σ反向传播解决了不可导的问题。三、PyTorch代码实现1. 定义VAE模型classVAE(nn.Module):def__init__(self,input_dim784,hidden_dim256,latent_dim20):super().__init__()self.encodernn.Sequential(nn.Linear(input_dim,hidden_dim),nn.ReLU(),nn.Linear(hidden_dim,hidden_dim),nn.ReLU())self.fc_munn.Linear(hidden_dim,latent_dim)self.fc_logvarnn.Linear(hidden_dim,latent_dim)self.decodernn.Sequential(nn.Linear(latent_dim,hidden_dim),nn.ReLU(),nn.Linear(hidden_dim,hidden_dim),nn.ReLU(),nn.Linear(hidden_dim,input_dim),nn.Sigmoid())defencode(self,x):hself.encoder(x)returnself.fc_mu(h),self.fc_logvar(h)defreparameterize(self,mu,log_var):stdtorch.exp(0.5*log_var)returnmutorch.randn_like(std)*stddefforward(self,x):mu,log_varself.encode(x)zself.reparameterize(mu,log_var)returnself.decode(z),mu,log_var2. 损失函数bce_lossnn.BCELoss(reductionsum)defloss_function(x_recon,x,mu,log_var):recon_lossbce_loss(x_recon,x)kl_loss-0.5*torch.sum(1log_var-mu.pow(2)-log_var.exp())returnrecon_losskl_loss3. 训练transformtransforms.Compose([transforms.ToTensor()])train_datasetdatasets.MNIST(root./data,trainTrue,downloadTrue,transformtransform)train_loaderDataLoader(train_dataset,batch_size128,shuffleTrue)devicetorch.device(cudaiftorch.cuda.is_available()elsecpu)modelVAE().to(device)optimizeroptim.Adam(model.parameters(),lr1e-3)forepochinrange(50):total_loss0fordata,_intrain_loader:datadata.view(-1,784).to(device)optimizer.zero_grad()x_recon,mu,log_varmodel(data)lossloss_function(x_recon,data,mu,log_var)loss.backward()total_lossloss.item()optimizer.step()print(fEpoch{epoch1}, Avg Loss:{total_loss/len(train_loader.dataset):.4f})4. 生成新数据model.eval()withtorch.no_grad():ztorch.randn(25,20).to(device)generated_imgsmodel.decode(z).cpu().numpy()⚠️注意本文仅为学习和理解算法进行 demo 代码实现线上和生产环境不建议使用。