在VAE中采用卷积形式，的pytorch代码

在PyTorch中实现基于卷积的变分自编码器（Convolutional Variational Autoencoder, ConvVAE），首先你需要安装PyTorch库。以下是一个简单的例子，展示了如何构建一个基础的卷积VAE模型：

import torch
import torch.nn as nn
from torch.autograd import Variable

class Encoder(nn.Module):
    def __init__(self, input_size, latent_dim, channels=3, kernel_size=4, stride=2, padding=1):
        super(Encoder, self).__init__()
        
        # 卷积层序列
        self.conv_layers = nn.Sequential(
            nn.Conv2d(channels, 64, kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        
        self.fc_mu = nn.Linear(128 * (input_size // 4) ** 2, latent_dim)
        self.fc_log_var = nn.Linear(128 * (input_size // 4) ** 2, latent_dim)

    def forward(self, x):
        out = self.conv_layers(x)
        out = out.view(out.size(0), -1)
        mu = self.fc_mu(out)
        log_var = self.fc_log_var(out)
        return mu, log_var

class Decoder(nn.Module):
    def __init__(self, latent_dim, output_size, channels=3, kernel_size=4, stride=2, padding=1):
        super(Decoder, self).__init__()
        
        self.deconv_layers = nn.Sequential(
            nn.ConvTranspose2d(latent_dim, 128, kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, kernel_size, stride=stride, padding=padding),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, channels, kernel_size, stride=1, padding=0)
        )

    def forward(self, z):
        out = z.view(z.size(0), -1, 1, 1)
        out = self.deconv_layers(out)
        return out

class ConvVAE(nn.Module):
    def __init__(self, input_size, latent_dim):
        super(ConvVAE, self).__init__()
        self.encoder = Encoder(input_size, latent_dim)
        self.decoder = Decoder(latent_dim, input_size)
    
    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)
        recon_x = self.decoder(z)
        return recon_x, mu, log_var

# 使用示例
input_shape = (1, 28, 28)  # 假设输入图片为灰度图像，大小为28x28
latent_dim = 10
model = ConvVAE(input_shape[1:], latent_dim)
x = torch.randn(1, *input_shape)
x = Variable(x)
recon_x, mu, log_var = model(x)

在这个例子中，我们定义了`Encoder`和`Decoder`模块，分别负责将输入数据编码成潜在向量和解码从潜在向量恢复原始数据。`ConvVAE`类组合了这两个组件，并包含了`reparameterize`函数用于采样。