变分自编码器 pytorch

变分自编码器（Variational Autoencoder，VAE）是一种生成模型，它可以学习数据的潜在分布，并用于生成新的数据。与传统的自编码器不同，VAE 引入了潜在变量（latent variable）的概念，使得模型更加灵活。

在 PyTorch 中，可以使用 `torch.nn` 模块来构建 VAE 模型。具体来说，需要定义编码器（encoder）、解码器（decoder）和潜在变量的分布。编码器将输入数据映射到潜在变量的分布上，解码器则将潜在变量映射回数据空间。训练时，需要最小化重构误差和 KL 散度，以使得模型能够学习到数据的潜在分布。

以下是一个简单的 VAE 实现示例：

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

class VAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(VAE, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.latent_dim = latent_dim
        
        # Encoder
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc21 = nn.Linear(hidden_dim, latent_dim)
        self.fc22 = nn.Linear(hidden_dim, latent_dim)
        
        # Decoder
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, input_dim)
        
    def encode(self, x):
        h = F.relu(self.fc1(x))
        mu = self.fc21(h)
        logvar = self.fc22(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mu + eps * std
        return z
    
    def decode(self, z):
        h = F.relu(self.fc3(z))
        x_hat = torch.sigmoid(self.fc4(h))
        return x_hat
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_hat = self.decode(z)
        return x_hat, mu, logvar