12、深入探索分层变分自编码器与扩散深度生成模型

深入探索分层变分自编码器与扩散深度生成模型

1. 自上而下变分自编码器（Top-down VAE）实现

我们来深入探讨自上而下变分自编码器（Top-down VAE）的实现。为了与前文描述相匹配，我们采用两级VAE。为使代码尽可能与数学表达式相似，我们使用单个类来实现，并采用重参数化技巧进行采样。在代码中，我们使用 $\log \sigma$ 而非 $\sigma$，使用 $\log \sigma + \log \sigma$ 而非 $\sigma \cdot \sigma$，因为 $e^{\log a + \log b} = e^{\log a} \cdot e^{\log b} = a \cdot b$。

以下是实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class HierarchicalVAE(nn.Module):
    def __init__(self, nn_r_1, nn_r_2, nn_delta_1, nn_delta_2,
                 nn_z_1, nn_x, num_vals=256, D=64, L=16, likelihood_type='categorical'):
        super(HierarchicalVAE, self).__init__()
        print('Hierachical VAE by JT.')
        # 自下而上路径
        self.nn_r_1 = nn_r_1
        self.nn