编写一个VAE pytorch版本

在PyTorch中实现Variational Autoencoder (VAE)是一个常见的深度学习任务，它是一种生成模型，用于学习数据的潜在分布并能够从这个分布中生成新的样本。以下是基本步骤：

# 导入必要的库
import torch
import torch.nn as nn
from torch.autograd import Variable

# 定义Encoder部分，这里包含一个输入层、隐藏层以及Q(z|x)网络（即编码器）
class Encoder(nn.Module):
    def __init__(self, input_dim, hid_dim, z_dim):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(input_dim, hid_dim)
        self.fc21 = nn.Linear(hid_dim, z_dim)
        self.fc22 = nn.Linear(hid_dim, z_dim)

    def forward(self, x):
        h1 = torch.relu(self.fc1(x))
        mean = self.fc21(h1)
        log_var = self.fc22(h1)
        return mean, log_var

# 定义Decoder部分，它接收z作为输入，并通过一系列线性和激活函数解码回原始空间
class Decoder(nn.Module):
    def __init__(self, z_dim, hid_dim, output_dim):
        super(Decoder, self).__init__()
        self.fc3 = nn.Linear(z_dim, hid_dim)
        self.fc4 = nn.Linear(hid_dim, hid_dim)
        self.fc5 = nn.Linear(hid_dim, output_dim)

    def forward(self, z):
        h2 = torch.relu(self.fc3(z))
        h3 = torch.relu(self.fc4(h2))
        out = torch.sigmoid(self.fc5(h3))  # 对于二进制数据，可以使用sigmoid；对于连续数据，可以使用tanh
        return out

# 创建完整的VAE类
class VAE(nn.Module):
    def __init__(self, encoder, decoder, z_dim):
        super(VAE, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.reparameterize = nn.ModuleList([nn.Linear(z_dim, z_dim), nn.Sigmoid()])  # 添加随机重参数化操作

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

    def forward(self, x):
        mean, log_var = self.encoder(x)
        z = self.reparametrize_(mean, log_var)
        recon_x = self.decoder(z)
        return recon_x, mean, log_var

# 使用例子
input_dim = ...  # 输入数据维度
hid_dim = ...   # 隐藏层维度
z_dim = ...     # 随机变量维度
model = VAE(Encoder(input_dim, hid_dim, z_dim), Decoder(z_dim, hid_dim, input_dim))