编写一个VAE pytorch

生成对抗网络（VAE，Variational Autoencoder）是一种深度学习模型，用于无监督的数据生成和潜在表示学习。在PyTorch中编写VAE，你需要准备一些基本组件，如编码器、解码器以及损失函数等。

以下是简单步骤和示例代码：

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

class Encoder(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(Encoder, self).__init__()
        # 你可以添加更多的层，这里是一个简单的线性层
        self.fc1 = nn.Linear(input_dim, 400)
        self.fc21 = nn.Linear(400, latent_dim) # mu
        self.fc22 = nn.Linear(400, latent_dim) # log(variance)

    def forward(self, x):
        h = F.relu(self.fc1(x))
        z_mu = self.fc21(h)
        z_log_var = self.fc22(h)
        return z_mu, z_log_var

class Decoder(nn.Module):
    def __init__(self, latent_dim, output_dim):
        super(Decoder, self).__init__()
        self.fc3 = nn.Linear(latent_dim, 400)
        self.fc4 = nn.Linear(400, output_dim)

    def forward(self, z):
        h = F.relu(self.fc3(z))
        x_recon = self.fc4(h)
        return x_recon

class VAE(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(VAE, self).__init__()
        self.encoder = Encoder(input_dim, latent_dim)
        self.decoder = Decoder(latent_dim, input_dim)

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, log_var = self.encoder(x)
        z = self.reparameterize(mu, log_var)
        x_recon = self.decoder(z)
        reconstruction_loss = F.mse_loss(x_recon, x, reduction='sum')
        kl_divergence = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        return reconstruction_loss, kl_divergence, x_recon

# 使用示例
vae = VAE(input_dim=784, latent_dim=20) # 假设输入数据是28x28灰度图片
x = torch.randn(100, 784)
reconstruction_loss, kl_div, recon_x = vae(x)