【深度学习】基于深度离散潜在变量模型的变分推理

本文链接：https://blog.csdn.net/MUKAMO/article/details/139993206

1.引言

1.1.讨论的目标

阅读并理解本文后，大家应能够：

掌握如何为具有离散潜在变量的模型设定参数
在可行的情况下，使用精确的对数似然函数来估计参数
利用神经变分推断方法来估计参数

1.2.导入相关软件包

# 导入PyTorch库，用于深度学习相关的操作。
import torch
# 导入NumPy库，用于进行科学计算。
import numpy as np
# 导入Python的random库，用于设置随机数生成器的种子。
import random
# 导入matplotlib.pyplot，用于绘图。
import matplotlib.pyplot as plt
# 导入PyTorch的神经网络模块。
import torch.nn as nn
# 导入PyTorch的函数库，用于一些常用的操作，如激活函数。
import torch.nn.functional as F
# 导入PyTorch的分布库，用于处理概率分布。
import torch.distributions as td

# 从itertools库导入chain函数，用于迭代多个可迭代对象。
from itertools import chain
# 从collections库导入defaultdict和OrderedDict，用于创建特殊类型的字典。
from collections import defaultdict, OrderedDict
# 从tqdm.auto导入tqdm，用于显示进度条。
from tqdm.auto import tqdm

# 从IPython.display导入set_matplotlib_formats，用于设置matplotlib的图形格式。
from IPython.display import set_matplotlib_formats
# 设置matplotlib图形的格式为SVG和PDF，SVG格式适合网页显示，PDF格式适合打印。
set_matplotlib_formats('svg', 'pdf') 
# 导入matplotlib配置，用于设置图形的全局参数。
import matplotlib
# 设置matplotlib图形的线宽。
matplotlib.rcParams['lines.linewidth'] = 2.0

# 确保matplotlib图形在Jupyter Notebook中以inline方式显示。
%matplotlib inline

# 定义一个函数seed_all，用于设置所有随机数生成器的种子，以确保实验的可复现性。
def seed_all(seed=42):
    # 设置NumPy的随机种子。
    np.random.seed(seed)
    # 设置Python内置random模块的种子。
    random.seed(seed)
    # 设置PyTorch的随机种子。
    torch.manual_seed(seed)

# 调用seed_all函数，设置默认的随机种子。
seed_all()

1.3.数据准备

为了本文的实验验证，我们将使用玩具图像数据集(FashionMNIST)。这些数据集包含固定维度的观测值，因此设计编码器和解码器相对简单。通过这种方式，我们可以将注意力集中在概率性质的核心部分。

# 从torchvision.datasets导入FashionMNIST数据集。
from torchvision.datasets import FashionMNIST
# 导入torchvision的transforms模块，用于数据预处理。
from torchvision import transforms
# 导入ToTensor转换，用于将PIL图像或Numpy数组转换为`FloatTensor`。
from torchvision.transforms import ToTensor
# 从torch.utils.data导入random_split和Dataset，用于数据集的划分和创建。
from torch.utils.data import random_split, Dataset
# 从torch.utils.data.dataloader导入DataLoader，用于创建数据加载器。
from torch.utils.data.dataloader import DataLoader
# 从torchvision.utils导入make_grid，用于将图像集合排列成一个网格。
from torchvision.utils import make_grid
# 导入PyTorch的优化器模块。
import torch.optim as opt

# 设置数据集的路径。
DATASET_PATH = "../../data"
# 根据是否有可用的GPU，设置运行设备。
my_device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")

# 定义一个二值化转换的Dataset类。
class Binarizer(Dataset):
    def __init__(self, ds, threshold=0.5):
        self._ds = ds  # 原始数据集
        self._threshold = threshold  # 二值化的阈值

    def __len__(self):
        # 返回数据集的大小。
        return len(self._ds)

    def __getitem__(self, idx):
        # 根据索引idx获取数据，并进行二值化处理。
        x, y = self._ds[idx]
        return (x >= self._threshold).float(), y

# 加载FashionMNIST数据集，并应用数据增强变换。
dataset = FashionMNIST(root=DATASET_PATH, train=True, download=True,
               transform=transforms.Compose([transforms.Resize(64), transforms.ToTensor()]))
# 打印出第一张图片的形状。
img_shape = dataset[0][0].shape
print("Shape of an image:", img_shape)

# 设置验证集的大小。
val_size = 1000
# 计算训练集的大小。
train_size = len(dataset) - val_size
# 划分训练集和验证集。
train_ds, val_ds = random_split(dataset, [train_size, val_size])

# 打印训练集和验证集的大小。
print(len(train_ds), len(val_ds))

# 设置是否将数据二值化的标志。
bin_data = True
if bin_data:
    # 如果需要，将训练集和验证集转换为二值化数据集。
    train_ds = Binarizer(train_ds)
    val_ds = Binarizer(val_ds)

# 设置批大小。
batch_size = 64
# 创建训练集的数据加载器。
train_loader = DataLoader(train_ds, batch_size, shuffle=True, num_workers=2, pin_memory=True)
# 创建验证集的数据加载器。
val_loader = DataLoader(val_ds, batch_size, num_workers=2, pin_memory=True)

# 遍历训练集的迭代器。
for images, y in train_loader:
    # 打印一批图像的形状。
    print('images.shape:', images.shape)
    # 创建一个图形，不显示坐标轴。
    plt.figure(figsize=(16,8))
    plt.axis('off')
    # 将图像集合排列成一个网格并显示。
    plt.imshow(make_grid(images, nrow=16).permute((1, 2, 0)))
    plt.show()
    # 展示一张图像后退出循环。
    break

2. 潜在变量模型

在本节中，我们将利用神经网络来定义一个涉及潜在变量的模型，也就是对一组随机变量的联合概率分布进行建模，其中部分变量可以被观察到，而另一部分则不能。

我们特别关注以下两种随机变量：

一个离散的潜在代码 $\in \mathcal{Z}$
一张图像 $\in \mathcal{X}$ ，它是实数空间 $RD\mathbb{R}^D$ 的一个子集

图像 $x$ 由 $C$ 个颜色通道、宽度 $W$ 和高度 $H$ 定义，因此 $X\mathcal{X}$ 属于 $RC×W×H\mathbb{R}^{C \times W \times H}$ 。由于我们设定了 $\times W \times H$ ， $X\mathcal{X}$ 是具有有限维度的，但这不是必须的（例如，在其他领域中， $X\mathcal{X}$ 可能是无限的，包含任意长度的所有句子）。我们可以将像素强度视为离散或连续变量，前提是我们为每种情况选择了恰当的概率质量函数或概率密度函数。

我们将探讨两种潜在代码类型：一种是分类代码 $\in \{1, \ldots, K\}$ ，另一种是组合代码 $\in \{0, 1\}^K$ 。在这两种情况下， $Z\mathcal{Z}$ 都是可数的，但这不是必须的（例如， $z$ 可以是自然数或任意长度的潜在序列）。

我们通过定义一个联合概率密度函数来指定 $Z\mathcal{Z}$ 和 $X\mathcal{X}$ 上的联合分布：

$pZX(z,x∣θ)=pZ(z∣θ)×pX∣Z(x∣z,θ)p_{ZX}(z, x|\theta) = p_Z(z|\theta) \times p_{X|Z}(x|z, \theta)$

这里的 $θ\theta$ 表示神经网络的参数，这些网络参数化了概率质量函数 $p_Z$ 和对于任何给定 $z$ 的条件概率密度函数 $p_{X|Z=z}$ 。

本文中的先验分布是固定的，但在其他情况下可能并非如此。我们没有其他预测变量作为条件，但在某些应用场景中可能会有（例如，在图像描述任务中，我们可能对给定图像 $x$ 的标题 $y$ 和潜在代码 $z$ 的联合分布感兴趣；在图像生成任务中，我们可能对给定标题 $y$ 的图像 $x$ 和潜在代码 $z$ 的联合分布感兴趣）。

2.1 先验网络

我们的讨论首先从定义和描述一个组件开始，这个组件用于参数化先验分布 $p_Z$ 。

先验网络是一种神经网络，其主要功能是为一批数据实例提供一个固定的先验分布的参数化表达。

class PriorNet(nn.Module):  
    """  
    先验网络：用于参数化先验分布的神经网络。  
  
    在本实验中，我们的先验是固定的，因此该神经网络的前向传播  
    只是返回一个具有给定batch_shape的固定先验。  
    """  
  
    def __init__(self, outcome_shape: tuple):  
        """  
        outcome_shape: 单个结果的形状。  
        如果你使用一个整数k，我们将会把它转换为(k,)的元组形式。  
        """  
        super().__init__()  
        if isinstance(outcome_shape, int):  
            # 如果outcome_shape是整数，则转换为元组形式  
            outcome_shape = (outcome_shape,)  
        self.outcome_shape = outcome_shape  
  
    def forward(self, batch_shape):  
        """  
        返回批次对应的先验分布对象。  
  
        Args:  
            batch_shape (tuple): 批次的形状。  
  
        Returns:  
            td object: 表示先验分布的torch.distributions对象（待实现）。  
  
        Raises:  
            NotImplementedError: 因为这个函数需要被实现。  
        """  
        raise NotImplementedError("需要实现此函数以返回先验分布对象！")  
  
# 注意：在实际实现中，您可能需要使用torch.distributions中的某个分布类来创建先验分布对象。  
# 例如，如果您想要一个正态分布的先验，您可以使用torch.distributions.Normal。

伯努利分布的乘积先验

当我们的潜在编码是一个 $K$ 维的比特向量时，每个比特都可以视为数据点的一个属性。对于这样的编码，我们为每一个坐标使用均匀分布的伯努利先验：

$p_Z(z) = \prod_{k=1}^K \text{Bernoulli}(z_k|0.5)$

这意味着每个比特都有0.5的概率是1，0.5的概率是0。

均匀独热类别先验

若潜在编码是从一个离散集合 $1,…,K}\{1, \ldots, K\}$ 中选取的一个类别，我们可以为该类别使用均匀的先验分布：

$p_Z(z) = \text{Categorical}(z|K^{-1} \mathbf{1}_K)$

其中， $1K\mathbf{1}_K$ 是一个 $K$ 维的全1向量。这个先验确保每个类别被选中的概率都是相等的，即 $1/ K$ 。

在实际应用中，我们可以使用“OneHotCategorical”分布来简化操作，它会自动将类别分布的样本转换为独热编码形式。

class BernoulliPriorNet(PriorNet):
    """伯努利先验网络，用于D维位向量z：
        p(z) = ∏ p(z[d] | 0.5)，其中p是伯努利分布。
    """
    
    def __init__(self, outcome_shape):
        super().__init__(outcome_shape)  # 调用基类的构造函数
        # 注册一个不需要梯度的缓冲区logits，初始化为0
        self.register_buffer("logits", torch.zeros(self.outcome_shape, requires_grad=False).detach())

    def forward(self, batch_shape):
        # 计算分布的形状
        shape = batch_shape + self.outcome_shape
        # 使用td.Independent来获得多变量抽样的概率质量函数（pmf）
        # 如果没有td.Independent，我们将得到多个pmf而不是一个多变量结果空间的pmf
        # td.Independent将最右边的维度解释为结果形状的一部分
        return td.Independent(td.Bernoulli(logits=self.logits.expand(shape)), len(self.outcome_shape))

class CategoricalPriorNet(PriorNet):
    """分类先验网络，用于z是K个类别集合中一个类别的一位有效编码：
        p(z) = OneHotCategorical(z | torch.ones(K) / K)
    """
    
    def __init__(self, outcome_shape):
        super().__init__(outcome_shape)  # 调用基类的构造函数
        self.register_buffer("logits", torch.zeros(self.outcome_shape, requires_grad=False).detach())

    def forward(self, batch_shape):
        shape = batch_shape + self.outcome_shape
        # OneHotCategorical是对Categorical的包装，
        # 在抽取Categorical样本后，td.OneHotCategorical使用onehot(sample, support_size)对其进行编码
        # 在这里我们不需要td.Independent，因为OneHotCategorical是对多变量抽样的概率分布
        # 这与伯努利先验的乘积不同
        return td.OneHotCategorical(logits=self.logits.expand(shape))

# 测试先验网络的函数
def test_priors(batch_size=3):
    prior_net = BernoulliPriorNet(7)
    print("\n伯努利先验")
    print(f" outcome_shape={
     
     prior_net.outcome_shape}")
    p = prior_net(batch_shape=(batch_size,))
    print(f" 分布: {
     
     p}")
    z = p.sample()  # 抽取一个样本
    print(f" 样本: {
     
     z}")
    print(f" 形状: sample={
     
     z.shape} log_prob={
     
     p.log_prob(z).shape}")

    prior_net = CategoricalPriorNet(7)
    print("\n分类先验")
    print(f" outcome_shape={
     
     prior_net.outcome_shape}")
    p = prior_net(batch_shape=(batch_size,))
    print(f" 分布: {
     
     p}")
    z = p.sample()
    print(f" 样本: {
     
     z}")
    print(f" 形状: sample={
     
     z.shape} log_prob={
     
     p.log_prob(z).shape}")

# 调用测试函数
test_priors()

伯努利先验
 outcome_shape=(7,)
分布: Independent(Bernoulli(logits: torch.Size([3, 7])), 1)
样本: tensor([[0., 0., 0., 1., 1., 1., 0.],
        [1., 0., 1., 0., 1., 0., 0.],
        [1., 0., 0., 1., 1., 1., 1.]])
形状: sample=torch.Size([3, 7]) log_prob=torch.Size([3])

分类先验
 outcome_shape=(7,)
分布: OneHotCategorical()
样本: tensor([[0., 0., 1., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0., 0.]])
形状: sample=torch.Size([3, 7]) log_prob=torch.Size([3])

2.2 条件概率分布

本节我们将通过神经网络来参数化条件概率分布（CPDs），具体做法是让神经网络学习并输出某个概率质量函数（PMF）或概率密度函数（PDF）的参数。这一步骤对于构建潜在变量模型中的 $p_{X|Z=z}$ 组件至关重要（并且，在后续的发展中，它也将对变分推断中的 $q_{Z|X=x}$ 组件有所帮助）。

我们的基本策略是，根据用户选择的输入，通过神经网络映射到由torch.distributions库支持的PMF或PDF的参数上。

class CPDNet(nn.Module):  
    """  
    CPDNet类：一个用于参数化条件概率分布的神经网络。  
  
    假设L是某个选定的分布，x ~ L 是具有outcome_shape形状的一个结果。  
  
    CPDNet的前向传播方法将多个输入映射到L的pmf/pdf的参数上，  
    并返回一个表示L的pmf/pdf的torch.distributions对象。  
    """  
  
    def __init__(self, outcome_shape):  
        """  
        初始化方法。  
  
        outcome_shape: 单个结果的形状。  
        如果传入一个整数k，则将其转换为元组形式(k,)。  
        """  
        super().__init__()  # 调用父类nn.Module的初始化方法  
        if isinstance(outcome_shape, int):  
            # 如果outcome_shape是整数，则转换为元组形式  
            outcome_shape = (outcome_shape,)  
        self.outcome_shape = outcome_shape  
  
    def forward(self, inputs):  
        """  
        前向传播方法。  
  
        根据输入`inputs`预测并返回一个torch.distributions对象。  
  
        inputs: 形状为batch_shape + (num_inputs,)的张量，  
                其中batch_shape是批次形状，num_inputs是输入特征的数量。  
  
        返回: 一个torch.distributions对象，表示根据输入预测的条件概率分布。  
  
        注意: 该方法需要被子类重写以实现具体的条件概率分布参数化逻辑。  
        """  
        raise NotImplementedError("请在此处实现具体的前向传播逻辑")  
  
# 注意：在实际使用中，子类应该实现forward方法以完成具体的条件概率分布参数化。