复杂一点的VAE模型，pytorch版本

在PyTorch中，复杂的Variational Autoencoder (VAE)模型通常会包含一些高级特性，比如深度结构（如多层感知机）、条件生成、或者结合其他技术（如自注意力机制）。例如，我们可以构建一个深度条件VAE (Conditional Variational Autoencoder，CVAE)，它接受额外的条件输入，并能根据这些条件生成特定样式的样本。

以下是一个简单的CVAE模型示例：

import torch
from torch import nn

class CVAE(nn.Module):
    def __init__(self, input_dim, hidden_dim, z_dim, num_layers=2):
        super(CVAE, self).__init__()
        self.encoder = nn.Sequential(
            # 输入层到隐藏层
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            # 隐藏层到中间层
            *([nn.Linear(hidden_dim, hidden_dim)] * num_layers),
            nn.ReLU(),
            # 中间层到z分布参数（均值和方差）
            nn.Linear(hidden_dim, z_dim * 2)
        )
        self.decoder = nn.Sequential(
            # z向量到隐藏层
            nn.Linear(z_dim, hidden_dim),
            nn.ReLU(),
            # 隐藏层到输出层
            *([nn.Linear(hidden_dim, hidden_dim)] * num_layers),
            nn.ReLU(),
            # 输出层，注意这里需要调整形状与输入匹配
            nn.Linear(hidden_dim, input_dim)
        )

    def encode(self, x):
        mean, log_var = self.encoder(x).chunk(2, dim=-1)
        return mean, log_var

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

    def decode(self, z):
        return self.decoder(z)

    def forward(self, x, cond_input=None):
        if cond_input is not None:
            x = torch.cat((x, cond_input), dim=1)
        mean, log_var = self.encode(x)
        z = self.reparameterize(mean, log_var)
        reconstructed_x = self.decode(z)
        return reconstructed_x, mean, log_var