03 - 变分推断基础¶

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

def kl_divergence_gaussian(mu1, logvar1, mu2, logvar2):
    """
    计算两个高斯分布之间的KL散度

    参数:
        mu1, logvar1: 分布q的均值和对数方差
        mu2, logvar2: 分布p的均值和对数方差

    返回:
        KL(q||p)
    """
    var1 = torch.exp(logvar1)
    var2 = torch.exp(logvar2)

    kl = 0.5 * (logvar2 - logvar1 + (var1 + (mu1 - mu2)**2) / var2 - 1)
    return kl.sum(dim=-1)

def compute_elbo(recon_x, x, mu, logvar):
    """
    计算VAE的ELBO

    参数:
        recon_x: 重构的图像
        x: 原始图像
        mu, logvar: 编码器输出的均值和对数方差

    返回:
        ELBO损失
    """
    # 重构损失（二元交叉熵或MSE）
    recon_loss = F.mse_loss(recon_x, x, reduction='sum') / x.shape[0]  # F.xxx PyTorch函数式API

    # KL散度
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.shape[0]

    # ELBO = - (重构损失 + KL散度)
    elbo = -(recon_loss + kl_loss)

    return elbo, recon_loss, kl_loss

# 测试
print("=" * 60)
print("ELBO计算测试")
print("=" * 60)

batch_size = 4
latent_dim = 10
image_size = 32

# 模拟数据
x = torch.randn(batch_size, 3, image_size, image_size)
recon_x = torch.randn(batch_size, 3, image_size, image_size)
mu = torch.randn(batch_size, latent_dim)
logvar = torch.randn(batch_size, latent_dim)

elbo, recon_loss, kl_loss = compute_elbo(recon_x, x, mu, logvar)

print(f"重构损失: {recon_loss.item():.4f}")  # 将单元素张量转为Python数值
print(f"KL散度: {kl_loss.item():.4f}")
print(f"ELBO: {elbo.item():.4f}")

class ReparameterizedGaussian:
    """
    可重参数化的高斯分布
    """
    def __init__(self, mu, logvar):
        self.mu = mu
        self.logvar = logvar
        self.std = torch.exp(0.5 * logvar)

    def sample(self, num_samples=None):
        """
        重参数化采样
        """
        if num_samples is None:
            eps = torch.randn_like(self.std)
        else:
            eps = torch.randn(num_samples, *self.std.shape)
        return self.mu + self.std * eps

    def rsample(self, num_samples=None):
        """
        可微采样（PyTorch风格）
        """
        return self.sample(num_samples)

    def log_prob(self, z):
        """
        计算对数概率
        """
        var = torch.exp(self.logvar)
        log_prob = -0.5 * (
            torch.log(2 * torch.pi * var) +
            (z - self.mu).pow(2) / var
        )
        return log_prob.sum(dim=-1)

    def kl_divergence(self, other_mu=0, other_logvar=0):
        """
        计算与标准高斯的KL散度
        """
        return kl_divergence_gaussian(
            self.mu, self.logvar,
            torch.zeros_like(self.mu),
            torch.zeros_like(self.logvar)
        )

# 测试重参数化
print("\n" + "=" * 60)
print("重参数化技巧测试")
print("=" * 60)

mu = torch.tensor([0.0, 1.0, -1.0])
logvar = torch.tensor([0.0, 0.5, -0.5])

dist = ReparameterizedGaussian(mu, logvar)

# 多次采样
samples = []
for _ in range(1000):
    z = dist.sample()
    samples.append(z)

samples = torch.stack(samples)  # torch.stack沿新维度拼接张量

print(f"理论均值: {mu}")
print(f"采样均值: {samples.mean(dim=0)}")
print(f"理论标准差: {torch.exp(0.5 * logvar)}")
print(f"采样标准差: {samples.std(dim=0)}")

# 验证梯度
mu_param = torch.tensor([0.0, 1.0], requires_grad=True)
logvar_param = torch.tensor([0.0, 0.0], requires_grad=True)

dist_grad = ReparameterizedGaussian(mu_param, logvar_param)
z = dist_grad.sample()
loss = z.pow(2).sum()
loss.backward()  # 反向传播计算梯度

print(f"\n梯度验证:")
print(f"mu梯度: {mu_param.grad}")
print(f"logvar梯度: {logvar_param.grad}")