生成一份cifar10的VAE代码

好的，以下是一个简单的 CIFAR-10 的 VAE 代码，使用 PyTorch 实现：

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import transforms
from torchvision.utils import save_image

# 定义超参数
batch_size = 128
lr = 1e-3
latent_size = 20
epochs = 50

# 加载数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 定义VAE模型
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        
        # 编码器
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(4096, latent_size * 2)
        )
        
        # 解码器
        self.decoder = nn.Sequential(
            nn.Linear(latent_size, 4096),
            nn.ReLU(),
            nn.Unflatten(1, (256, 4, 4)),
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 3, stride=1, padding=1),
            nn.Tanh()
        )

    def encode(self, x):
        h = self.encoder(x)
        mu = h[:, :latent_size]
        logvar = h[:, latent_size:]
        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):
        x_hat = self.decoder(z)
        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

# 定义损失函数和优化器
model = VAE()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss(reduction='sum')

# 训练模型
for epoch in range(1, epochs + 1):
    train_loss = 0

    for i, (x, _) in enumerate(train_dataloader):
        optimizer.zero_grad()

        x_hat, mu, logvar = model(x)
        loss_recon = criterion(x_hat, x)
        loss_kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        loss = loss_recon + loss_kld

        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch, epochs, train_loss / len(train_dataset)))

    # 每 10 个 epoch 保存一下重构后的图像
    if epoch % 10 == 0:
        with torch.no_grad():
            z = torch.randn(batch_size, latent_size)
            x_hat = model.decode(z)
            save_image(x_hat.view(-1, 3, 32, 32), 'samples/sample_{}.png'.format(epoch))

这个代码定义了一个 `VAE` 类，包含了编码器、解码器和前向传播的过程。在训练过程中，先将输入的图像通过编码器得到均值和方差，再通过重参数化技巧得到隐变量 $z$，最后通过解码器得到重构的图像。损失函数包括重构误差和 KL 散度两部分，使用 Adam 优化器进行模型优化。在训练过程中，每 10 个 epoch 保存一下重构后的图像。