• AE（自动编码器）与VAE（变分自动编码器）的区别和联系？

    他们各自的概念看以下链接就可以了：https://blog.csdn.net/weixin_43135178/category_11543123.html

     这里主要谈一下他们的区别？

    先说结论：

    • VAE是AE的升级版，VAE也可以被看作是一种特殊的AE
    • AE主要用于数据的压缩与还原，VAE主要用于生成。
    • AE是将数据映直接映射为数值code（确定的数值），而VAE是先将数据映射为分布，再从分布中采样得到数值code。
    • 损失函数和优化目标不同

    一、AE（Auto Encoder, 自动编码器）

    1、AE的结构

    如上图所示，自动编码器主要由两部分组成：编码器（Encoder）和解码器（Decoder）。编码器和解码器可以看作是两个函数，一个用于将高维输入（如图片）映射为低维编码（code），另一个用于将低维编码（code）映射为高维输出（如生成的图片）。这两个函数可以是任意形式，但在深度学习中，我们用神经网络去学习这两个函数。

    这时候我们只要拿出Decoder部分，随机生成一个code然后输入，就可以得到一张生成的图像。但实际上这样的生成效果并不好（下面解释原因），因此AE多用于数据压缩，而数据生成则使用下面所介绍的VAE更好。

    2、AE的缺陷

    由上面介绍可以看出，AE的Encoder是将图片映射成“数值编码”，Decoder是将“数值编码”映射成图片。这样存在的问题是，在训练过程中，随着不断降低输入图片与输出图片之间的误差，模型会过拟合，泛化性能不好。也就是说对于一个训练好的AE，输入某个图片，就只会将其编码为某个确定的code，输入某个确定的code就只会输出某个确定的图片，如果这个latent code来自于没见过的图片，那么生成的图片也不会好。下面举个例子来说明：

    假设我们训练好的AE将“新月”图片encode成code=1（这里假设code只有1维），将其decode能得到“新月”的图片；将“满月”encode成code=10，同样将其decode能得到“满月”图片。这时候如果我们给AE一个code=5，我们希望是能得到“半月”的图片，但由于之前训练时并没有将“半月”的图片编码，或者将一张非月亮的图片编码为5，那么我们就不太可能得到“半月”的图片。因此AE多用于数据的压缩和恢复，用于数据生成时效果并不理想。

    3、AE的代码实现

    3.1）AE encoder + decoder + AE的模型

    1. import torch
    2. from torch import nn
    3. from torch.autograd import Variable
    4.  
    5.  
    6. # Define the encoder and decoder networks
    7. class Encoder(nn.Module):
    8.     def __init__(self, input_dim, latent_dim):
    9.         super(Encoder, self).__init__()
    10.         self.fc1 = nn.Linear(input_dim, latent_dim)
    11.  
    12.     def forward(self, x):
    13.         x = torch.relu(self.fc1(x))
    14.         return x
    15.  
    16.  
    17. class Decoder(nn.Module):
    18.     def __init__(self, latent_dim, output_dim):
    19.         super(Decoder, self).__init__()
    20.         self.fc1 = nn.Linear(latent_dim, output_dim)
    21.  
    22.     def forward(self, x):
    23.         x = torch.sigmoid(self.fc1(x))
    24.         return x
    25.  
    26.  
    27. class Autoencoder(nn.Module):
    28.     def __init__(self, input_dim, latent_dim):
    29.         super(Autoencoder, self).__init__()
    30.         self.encoder = Encoder(input_dim, latent_dim)
    31.         self.decoder = Decoder(latent_dim, input_dim)
    32.  
    33.     def forward(self, x):
    34.         x = self.encoder(x)
    35.         x = self.decoder(x)
    36.         return x

    3.2）AE模型训练过程

    1. # Define the model, loss function, and optimizer
    2. input_dim = 784  # For MNIST images
    3. latent_dim = 32
    4. model = Autoencoder(input_dim, latent_dim)
    5. criterion = nn.BCELoss()
    6. optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    7.  
    8. # Training loop
    9. for epoch in range(num_epochs):
    10.     for data in dataloader:
    11.         img, _ = data
    12.         img = img.view(img.size(0), -1)
    13.         img = Variable(img)
    14.  
    15.         # Forward pass
    16.         output = model(img)
    17.         loss = criterion(output, img)
    18.  
    19.         # Backward pass and optimization
    20.         optimizer.zero_grad()
    21.         loss.backward()
    22.         optimizer.step()
    23.  
    24.     print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

    3.3）AE模型的推理过程

    1. # Evaluate the model
    2. model.eval()
    3. with torch.no_grad():
    4.     for data in dataloader:
    5.         img, _ = data
    6.         img = img.view(img.size(0), -1)
    7.         img = Variable(img)
    8.         output = model(img)

    3.4）AE怎么通过latent code生成新的图像

    1. # Initialize the autoencoder
    2. input_dim = 784  # Assuming input is a flattened 28x28 image
    3. latent_dim = 20  
    4. autoencoder = Autoencoder(input_dim, latent_dim)
    5. latent_vector = torch.randn(latent_dim)
    6.  
    7. # Generate a new image by passing the latent vector through the decoder
    8. with torch.no_grad(): 
    9.     generated_image = autoencoder.decoder(latent_vector)

    4、如何解决AE的问题呢？

    这时候我们转变思路，不将图片映射成“数值编码”，而将其映射成“分布”。还是刚刚的例子，我们将“新月”图片映射成μ=1的正态分布，那么就相当于在1附近加了噪声，此时不仅1表示“新月”，1附近的数值也表示“新月”，只是1的时候最像“新月”。将"满月"映射成μ=10的正态分布，10的附近也都表示“满月”。那么code=5时，就同时拥有了“新月”和“满月”的特点，那么这时候decode出来的大概率就是“半月”了。这就是VAE的思想。

    二、VAE（Variational Auto-Encoder, 变分自动编码器）

    1、VAE的结构

    vae和常规自编码器不同的地方在于它的encoder的output不是一个latent vector让decoder去做decode，而是从某个连续的分布里（常见的是高斯分布）采样得到一个随机数or随机向量，然后decode再去针对这个scaler做解码。

    不使用连续的AE这种分布，而是使用符合某种分布的VAE，这是因为：

    常规的ae的潜在空间的规律性是一个难点，它取决于初始空间中数据的分布、潜在空间的维度和编码器的架构等。因此，我们基本不可能就认为ae的latent vector的distribution和我们产生随机数的distribution是一个distribution（不可能），那就很尴尬了，假设latent vector的取值范围都在0～1之间，然后我们产生了一个包含了大量的负数的随机数让decoder去decode，那么decoder压根就decode不出什么正常的东西，training阶段压根就没见过嘛。想一想，潜在空间中的latent vector 的分布规律很难知道也是正常的,因为常规的自动编码器的任务中没有任何东西被训练来强制获得这样的规律(但是vae就会假设latent vector服从高斯分布)：自动编码器被训练成以尽可能少的损失进行编码和解码，压根就不care latent vector服从什么分布。那么自然,我们是不可能使用一个预定义的随机分布产生随机的input然后又期望decoder能够decode出有意义的东西的.

    既然我们不知道latent vector服从什么分布，我们就直接人为对其进行约束满足某种预定义的分布，这个预定义的分布和我们产生随机数的分布保持一致，不就完美解决问题了吗？

    所以通过VAE求出均值和方差，然后使用重参数化技巧在得到的这个分布中进行采样，就可以得到符合此分布的latent vector了。

    为什么使用重参数化？

    具体来说，在不使用重参数化的情况下，模型会直接从参数化的分布（例如，正态分布，由均值 μ 和方差 σ2 参数化）中采样，这使得梯度无法直接通过采样过程回传。重参数化技巧通过引入一个不依赖于模型参数的外部噪声源（通常是标准正态分布中抽取的），并对这个噪声进行变换（使用模型参数如均值和方差），来生成符合目标分布的样本。这样，模型的随机输出就可以表示为模型参数的确定性函数和一个随机噪声的组合。便可以完成梯度回传

    2、VAE的代码实现

    整体架构，VAE计算以下两方面之间的损失：

    1. 重构损失（Reconstruction Loss）：这一部分的损失计算的是输入数据与重构数据之间的差异。

    2. KL散度（Kullback-Leibler Divergence Loss）：这一部分的损失衡量的是学习到的潜在表示的分布与先验分布（通常假设为标准正态分布）之间的差异。KL散度是一种衡量两个概率分布相似度的指标，VAE通过最小化KL散度来确保学习到的潜在表示的分布尽可能接近先验分布。这有助于模型生成性能的提升，因为它约束了潜在空间的结构，使其更加规整，便于采样和推断。

    1）Encoder

    image --> 均值 + 标准差

    1. import torch
    2. from torch import nn
    3. from torch.nn import functional as F
    4.  
    5. # Encoder class definition
    6. class Encoder(nn.Module):
    7.     def __init__(self, input_dim, hidden_dim, latent_dim):
    8.         super(Encoder, self).__init__()
    9.         # 使用FC将输入变为隐藏层hidden_dim
    10.         self.fc1 = nn.Linear(input_dim, hidden_dim)
    11.         # Two fully connected layers to produce mean and log variance
    12.         # These will represent the latent space distribution parameters
    13.         self.fc21 = nn.Linear(hidden_dim, latent_dim) # 隐藏层hidden_dim --> 均值Mean μ
    14.         self.fc22 = nn.Linear(hidden_dim, latent_dim) # 隐藏层hidden_dim --> 标准差Log variance σ
    15.  
    16.     def forward(self, x):
    17.         # 使用RELU非线性变换，增加网络的表达能力
    18.         h1 = F.relu(self.fc1(x))
    19.         # Return the mean and log variance for the latent space
    20.         return self.fc21(h1), self.fc22(h1)

    2）Decoder

    1. # Decoder class definition
    2. class Decoder(nn.Module):
    3.     def __init__(self, latent_dim, hidden_dim, output_dim):
    4.         super(Decoder, self).__init__()
    5.         # latent_dim --> hidden_dim
    6.         self.fc3 = nn.Linear(latent_dim, hidden_dim)
    7.         # hidden_dim --> output_dim（输出的图像）
    8.         self.fc4 = nn.Linear(hidden_dim, output_dim)
    9.  
    10.     def forward(self, z):
    11.         h3 = F.relu(self.fc3(z))
    12.         return torch.sigmoid(self.fc4(h3))

    3）VAE

    重参数化技巧：

    这段代码对应的数学公式可以写作：

    z= σ⋅ϵ + μ

    其中：

    • z 是从潜在分布中采样得到的样本。
    • log var（log方差）:  log(σ2)
    • μ 是潜在分布的均值，对应代码中的 mu。
    • σ 是潜在分布的标准差 std ，通过 torch.exp(0.5*logvar) 计算得到，这里 logvar 是对数方差log(σ2)，因此 σ=exp(0.5*​log(σ2))。
    • ϵ 是从标准正态分布 N(0,1) 中采样得到的随机噪声，对应 torch.randn_like(std)。

    代码中的 eps.mul(std).add_(mu) 实现了上述公式的计算，即首先将随机噪声 ϵ 与标准差 σ 相乘，然后将结果加上均值 μ。这样，得到的 z 既包含了模型学习到的分布的特征（通过 μ 和 σ），同时也引入了必要的随机性（通过 ϵ），允许模型通过采样生成多样化的数据。

    1. # VAE class definition
    2. # Encode the input --> reparameterize --> decode
    3. class VAE(nn.Module):
    4.     def __init__(self, input_dim, hidden_dim, latent_dim):
    5.         super(VAE, self).__init__()
    6.         self.encoder = Encoder(input_dim, hidden_dim, latent_dim)
    7.         self.decoder = Decoder(latent_dim, hidden_dim, input_dim)
    8.  
    9.     def reparameterize(self, mu, logvar):
    10.         # Reparameterization trick to sample from the distribution represented by the mean and log variance
    11.         std = torch.exp(0.5*logvar)
    12.         eps = torch.randn_like(std)
    13.         return eps.mul(std).add_(mu)
    14.  
    15.     def forward(self, x):
    16.         mu, logvar = self.encoder(x.view(-1, input_dim))
    17.         z = self.reparameterize(mu, logvar)
    18.         return self.decoder(z), mu, logvar

    4）Loss

    1. # Loss function for VAE
    2. def vae_loss_function(recon_x, x, mu, logvar):
    3.     # Binary cross entropy between the target and the output
    4.     BCE = F.binary_cross_entropy(recon_x, x.view(-1, input_dim), reduction='sum')
    5.     # KL divergence loss ： 学习到的潜在表示的分布 <-->  先验分布(标准正态分布)
    6.     KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    7.     return BCE + KLD

    5）训练过程

    1. # Hyperparameters
    2. input_dim = 784 # Assuming input is a flattened 28x28 image (e.g., from MNIST)
    3. hidden_dim = 400
    4. latent_dim = 20
    5. epochs = 10
    6. learning_rate = 1e-3
    7.  
    8. # Initialize VAE
    9. vae = VAE(input_dim, hidden_dim, latent_dim)
    10. optimizer = torch.optim.Adam(vae.parameters(), lr=learning_rate)
    11.  
    12. # Training process function
    13. for epoch in range(epochs):
    14.     vae.train()  # Set the model to training mode
    15.     train_loss = 0
    16.     for batch_idx, (data, _) in enumerate(data_loader):
    17.         optimizer.zero_grad()  # Zero the gradients
    18.         recon_batch, mu, logvar = vae(data)  # Forward pass through VAE
    19.         loss = vae_loss_function(recon_batch, data, mu, logvar)  # Compute the loss
    20.         loss.backward()  # Backpropagate the loss
    21.         train_loss += loss.item()
    22.         optimizer.step()

    三、VQ-VAE

    之所以叫做VQ（向量量化），主要是因为将连续潜在空间的点映射到最近的一组离散的向量（即码本中的向量）上

    VQ-VAE的全称是Vector Quantized-Variational AutoEncoder，即向量量化变分自编码器。这是一种结合了变分自编码器（VAE）和向量量化（VQ）的深度学习模型，主要用于高效地学习数据的潜在表示。VQ-VAE通过将连续的潜在表示空间离散化来改进传统VAE模型。向量量化的过程实质上是将连续潜在空间的点映射到最近的一组离散的向量（即码本中的向量）上，这有助于模型捕捉和表示更加丰富和复杂的数据分布，由于维护了一个codebook，编码范围更加可控，VQVAE相对于VAE（VAE的隐变量 z 的每一维都是一个连续的值, 而VQ-VAE最大的特点就是, z 的每一维都是离散的整数。），可以生成更大更高清的图片(这也为后续DALLE和VQGAN的出现做了铺垫)。【原文中说的是避免了“后验坍塌”的问题】

    1、算法步骤：

    1. 通过Encoder学习出中间编码 Ze(x)【绿色】
    2. 事先定义好codebook，它有N个e组成【紫色】
    3. 然后通过最邻近搜索与中间编码Ze(x)最相似（接近）的codebook中K个向量之一，并记住这个向量的index【青色】
    4. 根据得到的所有index去映射对应的codebook中的vector，得到输入图像对应的特征表征Zq(x)【紫色】
    5. 然后通过Decoder对Zq(x)进行重建

    另外由于最邻近搜索使用argmax来找codebook中的索引位置，导致不可导问题，VQVAE通过stop gradient操作来避免最邻近搜索的不可导问题，也就是通过stop gradient操作，将decoder输入的梯度复制到encoder的输出上【红色的线】。

    2、一些问题

    A. 为什么要进行向量量化？（为什么要将 z 离散化）？

    1. 离散的表示通常更适合于捕捉数据中的类别性质，如不同种类的对象、语音或文本数据的不同模式等。

    2. 离散的潜在表示有助于模型生成更加清晰的输出。在连续潜在空间中，模型可能在生成新样本时产生模糊的结果，特别是在空间的某些区域中。而离散潜在空间能够降低这种模糊性，提高生成样本的质量。

    3. 增强模型的解释性，相比于连续潜在空间，离散潜在空间可以为每个离散的潜在变量赋予更明确的语义解释。例如，在处理图像的任务中，不同的离散潜在变量可能对应于不同的视觉特征或对象类别，这使得模型的行为和学习到的表示更易于理解和解释。

    5. 缓解潜在空间的过度平滑问题，VAE有时会遇到潜在空间的"过度平滑"问题，即潜在空间中不同区域之间的过渡太平滑，导致生成的样本缺乏多样性或区分度不够（容易模型崩塌）。通过引入离散潜在空间，VQ-VAE可以缓解这个问题，因为离散空间天然具有区分不同区域的能力。

    B. 如何将 z 离散化？

    1）构建codebook进行VQ 

    将 z 离散化的关键就是VQ, 即vector quatization.

    简单来说, 就是要先有一个codebook, 这个codebook是一个embedding table，然后再利用均匀分布对权重初始化

    1. self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
    2. self._embedding.weight.data.uniform_(-1 / self._num_embeddings, 1 / self._num_embeddings)

    2）找到与图像embedding（flat_input）最近的codebook中的embedding 

    我们在这个codebook 中找到和 vector 最接近(比如欧氏距离最近)的一个embedding, 用这个embedding的index来代表这个vector.

    ∑ [flat_input(16384, 64) - self._embedding.weight(512,64)]^2

    1. # Calculate the Z_e(x) and e distances
    2. # 这里使用欧几里得距离的平方求距离
    3. distances = (torch.sum(flat_input ** 2, dim=1, keepdim=True)
    4.              + torch.sum(self._embedding.weight ** 2, dim=1)
    5.              - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
    6.  
    7. # Encoding
    8. encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)

    3）怎么解决的梯度截断（不可导）问题？

    另外由于最邻近搜索使用argmin来找codebook中的索引位置，导致不可导问题，VQVAE通过stop gradient操作来避免最邻近搜索的不可导问题，也就是通过stop gradient操作，将decoder输入（quantize）的梯度复制到encoder的输出上（input）。

    1. quantize = input + (quantize - input).detach()
    2. # 正向传播和往常一样
    3.  
    4. # 反向传播时，detach()这部分梯度为0，quantize和input的梯度相同
    5. # 即实现将quantize复制给input

    VQVAE相比于VAE最大的不同是，直接找每个属性的离散值，通过类似于查表的

    直接按照下面这样，而不进行 (quantized - inputs).detach() 是不可以的，因为这样会让模型退化到一个VAE的样子，因为你是直接将输出复制给输入的。具体没搞清楚，反正人家是这么干的...

    quantize = input

    3、VQVAE的损失

    与VAE的不同：去掉VAE的KL loss，增加了两项loss

    1）重构损失（reconstruction loss）：

    • 目标：衡量重构数据和原始数据之间的相似度。
    • 计算：通常使用均方误差（MSE）或交叉熵损失来计算重构图像和原始图像之间的差异。

    2）代码本损失（codebook loss）：

    代码本损失关注于更新码本向量，使其更好地代表输入数据的连续潜在表示。

    假设Ze​(x)是编码器对输入x的连续潜在表示，e是选取的最接近的码本向量。

    3）提交损失（Commitment Loss）：

    提交损失则确保编码器的输出不会偏离它选择的码本向量太远，从而保证训练过程的稳定性。

    计算模型编码器输出的连续潜在表示和量化后的表示（即选取的码本向量）之间的距离。这有助于稳定训练，确保编码器输出与码本的选择保持一致，也通常使用均方误差（MSE）计算。

    假设Ze​(x)是编码器对输入x的连续潜在表示，e是选取的最接近的码本向量。

    4、代码示例

    需要确保经过encoder的图像的通道数D（此时经过encoder后的图像“通道数”不一定再是3了，可能会更大，例如64，我们这里只是形象的把它叫做“通道数”罢了）与codebook中的向量维度是相同的。

    整体步骤

    具体实现步骤如下：

    完整代码如下：

    1. from __future__ import print_function
    2.  
    3.  
    4. import matplotlib.pyplot as plt
    5. import numpy as np
    6. from scipy.signal import savgol_filter
    7.  
    8.  
    9. from six.moves import xrange
    10.  
    11. import umap
    12.  
    13. import torch
    14. import torch.nn as nn
    15. import torch.nn.functional as F
    16. from torch.utils.data import DataLoader
    17. import torch.optim as optim
    18.  
    19. import torchvision.datasets as datasets
    20. import torchvision.transforms as transforms
    21. from torchvision.utils import make_grid
    22.  
    23.  
    24. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    25.  
    26.  
    27. # ===================================================== Load Data =====================================================
    28. training_data = datasets.CIFAR10(root="data", train=True, download=True,
    29.                                   transform=transforms.Compose([
    30.                                       transforms.ToTensor(),
    31.                                       transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
    32.                                   ]))
    33.  
    34. validation_data = datasets.CIFAR10(root="data", train=False, download=True,
    35.                                   transform=transforms.Compose([
    36.                                       transforms.ToTensor(),
    37.                                       transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))
    38.                                   ]))
    39. # 计算 方差
    40. data_variance = np.var(training_data.data / 255.0)
    41.  
    42.  
    43.  
    44. # ===================================================== Vector Quantizer Layer =====================================================
    45. """
    46. This layer takes a tensor to be quantized. 
    47. The channel dimension will be used as the space in which to quantize. All other dimensions will be flattened and will be seen as different examples to quantize.
    48. The output tensor will have the same shape as the input.
    49. As an example for a BCHW tensor of shape [16, 64, 32, 32], we will first convert it to an BHWC tensor of shape [16, 32, 32, 64] and then reshape it into [16384, 64] and all 16384 vectors of size 64 will be quantized independently. 
    50. In otherwords, the channels are used as the space in which to quantize. All other dimensions will be flattened and be seen as different examples to quantize, 16384 in this case.
    51. """
    52.  
    53. class VectorQuantizer(nn.Module):
    54.     def __init__(self, num_embeddings, embedding_dim, commitment_cost):
    55.         '''
    57.         :param num_embeddings: codebook的大小
    58.         :param embedding_dim: codebook中每个vector的维度
    59.         :param commitment_cost: commit loss的β
    60.         '''
    61.  
    62.         super(VectorQuantizer, self).__init__()
    63.  
    64.         self._embedding_dim = embedding_dim
    65.         self._num_embeddings = num_embeddings
    66.         # 构建一个codebook，用均匀分布对codebook的权重进行初始化
    67.         self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
    68.         self._embedding.weight.data.uniform_(-1 / self._num_embeddings, 1 / self._num_embeddings)
    69.         self._commitment_cost = commitment_cost
    70.  
    71.     def forward(self, inputs):
    72.  
    73.         # convert inputs(encoder's output) from BCHW -> BHWC
    74.         inputs = inputs.permute(0, 2, 3, 1).contiguous()
    75.         input_shape = inputs.shape
    76.  
    77.         # Flatten input
    78.         flat_input = inputs.view(-1, self._embedding_dim)
    79.  
    80.         # Calculate the Z_e(x) and e distances
    81.         # 这里使用欧几里得距离的平方求距离distance(16384,512)
    82.         distances = (torch.sum(flat_input ** 2, dim=1, keepdim=True)
    83.                      + torch.sum(self._embedding.weight ** 2, dim=1)
    84.                      - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
    85.  
    86.         # 通过distance得到距离最近的index
    87.         encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
    88.         # 转为one-hot格式
    89.         encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
    90.         encodings.scatter_(1, encoding_indices, 1)
    91.  
    92.         # Quantize and unflatten：得到最近邻的Embedding Vector
    93.         quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
    94.  
    95.         # Loss
    96.         # commit loss
    97.         e_latent_loss = F.mse_loss(quantized.detach(), inputs)
    98.         # codebook loss
    99.         q_latent_loss = F.mse_loss(quantized, inputs.detach())
    100.         loss = q_latent_loss + self._commitment_cost * e_latent_loss
    101.         # trick（梯度复制），通过添加一个常数让编码器和解码器连续可导
    102.         quantized = inputs + (quantized - inputs).detach()
    103.         # 利用困惑度监测分布，困惑度越大，信息熵也就越大，分布就没有这么均匀
    104.         avg_probs = torch.mean(encodings, dim=0)
    105.         perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
    106.  
    107.         # convert quantized from BHWC -> BCHW
    108.         return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings
    109.  
    110.  
    111. """
    112. We will also implement a slightly modified version which will use exponential moving averages to update the embedding vectors instead of an auxillary loss.
    113. This has the advantage that the embedding updates are independent of the choice of optimizer for the encoder, decoder and other parts of the architecture. 
    114. For most experiments the EMA version trains faster than the non-EMA version.
    115. """
    116.  
    117. class VectorQuantizerEMA(nn.Module):
    118.     def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5):
    119.         super(VectorQuantizerEMA, self).__init__()
    120.  
    121.         self._embedding_dim = embedding_dim
    122.         self._num_embeddings = num_embeddings
    123.  
    124.         self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)
    125.         self._embedding.weight.data.normal_()
    126.         self._commitment_cost = commitment_cost
    127.  
    128.         self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings))
    129.         self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim))
    130.         self._ema_w.data.normal_()
    131.  
    132.         self._decay = decay
    133.         self._epsilon = epsilon
    134.  
    135.     def forward(self, inputs):
    136.         # convert inputs from BCHW -> BHWC
    137.         inputs = inputs.permute(0, 2, 3, 1).contiguous()
    138.         input_shape = inputs.shape
    139.  
    140.         # Flatten input
    141.         flat_input = inputs.view(-1, self._embedding_dim)
    142.  
    143.         # Calculate distances
    144.         distances = (torch.sum(flat_input ** 2, dim=1, keepdim=True)
    145.                      + torch.sum(self._embedding.weight ** 2, dim=1)
    146.                      - 2 * torch.matmul(flat_input, self._embedding.weight.t()))
    147.  
    148.         # Encoding
    149.         encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
    150.         encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
    151.         encodings.scatter_(1, encoding_indices, 1)
    152.  
    153.         # Quantize and unflatten
    154.         quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
    155.  
    156.         # Use EMA to update the embedding vectors
    157.         if self.training:
    158.             self._ema_cluster_size = self._ema_cluster_size * self._decay + \
    159.                                      (1 - self._decay) * torch.sum(encodings, 0)
    160.  
    161.             # Laplace smoothing of the cluster size
    162.             n = torch.sum(self._ema_cluster_size.data)
    163.             self._ema_cluster_size = (
    164.                     (self._ema_cluster_size + self._epsilon)
    165.                     / (n + self._num_embeddings * self._epsilon) * n)
    166.  
    167.             dw = torch.matmul(encodings.t(), flat_input)
    168.             self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw)
    169.  
    170.             self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1))
    171.  
    172.         # Loss
    173.         e_latent_loss = F.mse_loss(quantized.detach(), inputs)
    174.         loss = self._commitment_cost * e_latent_loss
    175.  
    176.         # Straight Through Estimator
    177.         quantized = inputs + (quantized - inputs).detach()
    178.         avg_probs = torch.mean(encodings, dim=0)
    179.         perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
    180.  
    181.         # convert quantized from BHWC -> BCHW
    182.         return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings
    183.  
    184.  
    185. # ===================================================== Encoder & Decoder Architecture =====================================================
    186. class Residual(nn.Module):
    187.     def __init__(self, in_channels, num_hiddens, num_residual_hiddens):
    188.         super(Residual, self).__init__()
    189.         self._block = nn.Sequential(
    190.             nn.ReLU(True),
    191.             nn.Conv2d(in_channels=in_channels,
    192.                       out_channels=num_residual_hiddens,
    193.                       kernel_size=3, stride=1, padding=1, bias=False),
    194.             nn.ReLU(True),
    195.             nn.Conv2d(in_channels=num_residual_hiddens,
    196.                       out_channels=num_hiddens,
    197.                       kernel_size=1, stride=1, bias=False)
    198.         )
    199.  
    200.     def forward(self, x):
    201.         return x + self._block(x)
    202.  
    203.  
    204. class ResidualStack(nn.Module):
    205.     def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
    206.         super(ResidualStack, self).__init__()
    207.         self._num_residual_layers = num_residual_layers
    208.         self._layers = nn.ModuleList([Residual(in_channels, num_hiddens, num_residual_hiddens)
    209.                                       for _ in range(self._num_residual_layers)])
    210.  
    211.     def forward(self, x):
    212.         for i in range(self._num_residual_layers):
    213.             x = self._layers[i](x)
    214.         return F.relu(x)
    215.  
    216.  
    217. class Encoder(nn.Module):
    218.     def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
    219.         super(Encoder, self).__init__()
    220.  
    221.         self._conv_1 = nn.Conv2d(in_channels=in_channels,
    222.                                  out_channels=num_hiddens // 2,
    223.                                  kernel_size=4,
    224.                                  stride=2, padding=1)
    225.         self._conv_2 = nn.Conv2d(in_channels=num_hiddens // 2,
    226.                                  out_channels=num_hiddens,
    227.                                  kernel_size=4,
    228.                                  stride=2, padding=1)
    229.         self._conv_3 = nn.Conv2d(in_channels=num_hiddens,
    230.                                  out_channels=num_hiddens,
    231.                                  kernel_size=3,
    232.                                  stride=1, padding=1)
    233.         self._residual_stack = ResidualStack(in_channels=num_hiddens,
    234.                                              num_hiddens=num_hiddens,
    235.                                              num_residual_layers=num_residual_layers,
    236.                                              num_residual_hiddens=num_residual_hiddens)
    237.  
    238.     def forward(self, inputs):
    239.         x = self._conv_1(inputs)
    240.         x = F.relu(x)
    241.  
    242.         x = self._conv_2(x)
    243.         x = F.relu(x)
    244.  
    245.         x = self._conv_3(x)
    246.         return self._residual_stack(x)
    247.  
    248.  
    249. class Decoder(nn.Module):
    250.     def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):
    251.         super(Decoder, self).__init__()
    252.  
    253.         self._conv_1 = nn.Conv2d(in_channels=in_channels,
    254.                                  out_channels=num_hiddens,
    255.                                  kernel_size=3,
    256.                                  stride=1, padding=1)
    257.  
    258.         self._residual_stack = ResidualStack(in_channels=num_hiddens,
    259.                                              num_hiddens=num_hiddens,
    260.                                              num_residual_layers=num_residual_layers,
    261.                                              num_residual_hiddens=num_residual_hiddens)
    262.  
    263.         self._conv_trans_1 = nn.ConvTranspose2d(in_channels=num_hiddens,
    264.                                                 out_channels=num_hiddens // 2,
    265.                                                 kernel_size=4,
    266.                                                 stride=2, padding=1)
    267.  
    268.         self._conv_trans_2 = nn.ConvTranspose2d(in_channels=num_hiddens // 2,
    269.                                                 out_channels=3,
    270.                                                 kernel_size=4,
    271.                                                 stride=2, padding=1)
    272.  
    273.     def forward(self, inputs):
    274.         x = self._conv_1(inputs)
    275.  
    276.         x = self._residual_stack(x)
    277.  
    278.         x = self._conv_trans_1(x)
    279.         x = F.relu(x)
    280.  
    281.         return self._conv_trans_2(x)
    282.  
    283.  
    284. # ===================================================== Train =====================================================
    285. batch_size = 256
    286. num_training_updates = 15000
    287.  
    288. num_hiddens = 128
    289. num_residual_hiddens = 32
    290. num_residual_layers = 2
    291.  
    292. # codebook的维度
    293. embedding_dim = 64
    294. num_embeddings = 512
    295.  
    296. commitment_cost = 0.25
    297.  
    298. decay = 0.99
    299. decay = 0
    300. learning_rate = 1e-3
    301.  
    302. training_loader = DataLoader(training_data,
    303.                              batch_size=batch_size,
    304.                              shuffle=True,
    305.                              pin_memory=True)
    306. validation_loader = DataLoader(validation_data,
    307.                                batch_size=32,
    308.                                shuffle=True,
    309.                                pin_memory=True)
    310.  
    311.  
    312. class Model(nn.Module):
    313.     def __init__(self, num_hiddens, num_residual_layers, num_residual_hiddens, num_embeddings, embedding_dim, commitment_cost, decay=0):
    314.         super(Model, self).__init__()
    315.  
    316.         self._encoder = Encoder(3, num_hiddens, num_residual_layers, num_residual_hiddens)
    317.         # 对encoder的输出进行后处理，得到与embedding table一样大小的维度
    318.         self._pre_vq_conv = nn.Conv2d(in_channels=num_hiddens,
    319.                                       out_channels=embedding_dim,
    320.                                       kernel_size=1,
    321.                                       stride=1)
    322.         if decay > 0.0:
    323.             self._vq_vae = VectorQuantizerEMA(num_embeddings, embedding_dim, commitment_cost, decay)
    324.         else:
    325.             self._vq_vae = VectorQuantizer(num_embeddings, embedding_dim, commitment_cost)
    326.  
    327.         self._decoder = Decoder(embedding_dim, num_hiddens, num_residual_layers, num_residual_hiddens)
    328.  
    329.     def forward(self, x):
    330.         z = self._encoder(x)
    331.         z = self._pre_vq_conv(z)
    332.         loss, quantized, perplexity, _ = self._vq_vae(z)
    333.         x_recon = self._decoder(quantized)
    334.  
    335.         return loss, x_recon, perplexity
    336.  
    337.  
    338. model = Model(num_hiddens, num_residual_layers, num_residual_hiddens, num_embeddings, embedding_dim, commitment_cost, decay).to(device)
    339. optimizer = optim.Adam(model.parameters(), lr=learning_rate, amsgrad=False)
    340. model.train()
    341. train_res_recon_error = []
    342. train_res_perplexity = []
    343.  
    344. for i in xrange(num_training_updates):
    345.     (data, _) = next(iter(training_loader))
    346.     data = data.to(device)
    347.     optimizer.zero_grad()
    348.  
    349.     vq_loss, data_recon, perplexity = model(data)
    350.     recon_error = F.mse_loss(data_recon, data) / data_variance
    351.     loss = recon_error + vq_loss
    352.     loss.backward()
    353.  
    354.     optimizer.step()
    355.  
    356.     train_res_recon_error.append(recon_error.item())
    357.     train_res_perplexity.append(perplexity.item())
    358.  
    359.     if (i + 1) % 100 == 0:
    360.         print('%d iterations' % (i + 1))
    361.         print('recon_error: %.3f' % np.mean(train_res_recon_error[-100:]))
    362.         print('perplexity: %.3f' % np.mean(train_res_perplexity[-100:]))
    363.         print()
    364.

    pytorch-vq-vae/vq-vae.ipynb at master · zalandoresearch/pytorch-vq-vae · GitHub

    四、总结

    • AE主要用于数据的压缩与还原，在生成数据上使用VAE。
    • AE是将数据映直接映射为数值code，而VAE是先将数据映射为分布，再从分布中采样得到数值code。
    • VQ-VAE是将中间编码映射为codebook中K个向量之一，然后通过Decoder对latent code进行重建

    因此AutoEncoder、VAE和VQ-VAE可以统一为latent code的概率分布设计不一样，AutoEncoder通过网络学习得到任意概率分布，VAE设计为正态分布，VQVAE设计为codebook的离散分布。总之，AutoEncoder的重构思想就是用低纬度的latent code分布来表达高纬度的数据分布，VAE和VQVAE的重构思想是通过设计latent code的分布形式，进而控制图片生成的过程。

    漫谈VAE和VQVAE，从连续分布到离散分布 - 知乎

  • 相关阅读:
    数据通讯基础
    Gateway路由的配置方式
    （八)RabbitMQ发布确认
    Servlet | 域对象、request对象其它常用的方法
    3如何搭建组件库的样式工程之button-scss
    如何实现 System.out.println("a") 显示 b
    怎么在国家级媒体网站投稿发布通稿
    Linux目录结构
    如何快速高效压缩图片？
    AWS SAP-C02 考试指南
  • 原文地址：https://blog.csdn.net/weixin_43135178/article/details/130592568