Registro simple VAE (Codificador automático variacional)

prefacio

Lo encuentro a menudo, pero lo olvido cada vez que hago un pequeño suplemento, por lo que también podría abrir un artículo y registrarlo lentamente.
VAE -> VQVAE, agregando principalmente cuantificación vectorial

Sobre VQ-VAE, puede consultar este VQ-VAE del modelo generado

Este artículo se actualizará continuamente…

Este artículo está bien escrito, mírelo bien cuando tenga tiempo Autoencoder variacional VAE: Esto es lo que es | Código fuente abierto adjunto , un artículo del Sr. Su Jianlin.

teoría

fase de entrenamiento

Para decirlo sin rodeos, ingrese una muestra (imagen), y después de que el codificador extraiga características de esta muestra, aprende dos cantidades, una es el valor medio y la otra es la varianza.
característica = codificador (img)
mu, var = w_mu (característica), w_var (característica)
Luego, la variable oculta z se puede muestrear de acuerdo con la media y la varianza.
eps = antorcha.rand_like(mu)
z = mu + (var) ** 0.5 * eps

Luego, de acuerdo con esta variable oculta z, podemos decodificar la imagen.
img_generar = decodificador (z)

Si la dimensión es incorrecta, puede agregar una capa completamente conectada para mapear la dimensión.

etapa de pronóstico

Primero genere una variable oculta aleatoria z, preste atención a la dimensión de z
y luego introdúzcala en el decodificador.

el código

La fuente del siguiente código: https://zhuanlan.zhihu.com/p/151587288
El código de entrenamiento se puede ver en https://shenxiaohai.me/2018/10/20/pytorch-tutorial-advanced-02/
es hacer un crossEntropyloss con la verdad fundamental, más un KLloss

class VAE(nn.Module):
    def __init__(self, latent_dim):
        super().__init__()
        
        self.encoder = nn.Sequential(nn.Linear(28 * 28, 256),
                                     nn.ReLU(),
                                     nn.Linear(256, 128))
        
        self.mu     = nn.Linear(128, latent_dim)
        self.logvar = nn.Linear(128, latent_dim)
        
        self.latent_mapping = nn.Linear(latent_dim, 128)
        
        self.decoder = nn.Sequential(nn.Linear(128, 256),
                                     nn.ReLU(),
                                     nn.Linear(256, 28 * 28))
        
        
    def encode(self, x):
        x = x.view(x.size(0), -1)
        encoder = self.encoder(x)
        mu, logvar = self.mu(encoder), self.logvar(encoder)
        return mu, logvar
        
    def sample_z(self, mu, logvar):
        eps = torch.rand_like(mu)
        return mu + eps * torch.exp(0.5 * logvar)
    
    def decode(self, z,x):
        latent_z = self.latent_mapping(z)
        out = self.decoder(latent_z)
        reshaped_out = torch.sigmoid(out).view(x.shape[0],1, 28,28)
        return reshaped_out
        
    def forward(self, x):
        
        mu, logvar = self.encode(x)
        z = self.sample_z(mu, logvar)
        output = self.decode(z,x)
        
        return output

# 创建优化器
num_epochs = 10
learning_rate = 1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    for i, (x, _) in enumerate(data_loader):
        # 获取样本，并前向传播
        x = x.to(device).view(-1, 28 * 28)
        x_predict = model(x)
        
        # 计算重构损失和KL散度（KL散度用于衡量两种分布的相似程度）
        # KL散度的计算可以参考论文或者文章开头的链接
        reconst_loss = F.binary_cross_entropy(x_predict, x, size_average=False)
        kl_div = - 0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        
        # 反向传播和优化
        loss = reconst_loss + kl_div
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()