pytorch实现VAE

PyTorch可以用于实现变分自编码器(VAE)。以下是一个简单的PyTorch代码示例，演示了如何实现VAE模型：

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

# 定义编码器网络
class Encoder(nn.Module):
    def __init__(self, latent_dim):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2_mean = nn.Linear(256, latent_dim)
        self.fc2_logvar = nn.Linear(256, latent_dim)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        mean = self.fc2_mean(x)
        logvar = self.fc2_logvar(x)
        return mean, logvar

# 定义解码器网络
class Decoder(nn.Module):
    def __init__(self, latent_dim):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(latent_dim, 256)
        self.fc2 = nn.Linear(256, 784)
    
    def forward(self, z):
        z = F.relu(self.fc1(z))
        x = torch.sigmoid(self.fc2(z))
        return x

# 定义VAE模型
class VAE(nn.Module):
    def __init__(self, latent_dim):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim
        self.encoder = Encoder(latent_dim)
        self.decoder = Decoder(latent_dim)
    
    def reparameterize(self, mean, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        z = mean + eps * std
        return z
    
    def forward(self, x):
        mean, logvar = self.encoder(x.view(-1, 784))
        z = self.reparameterize(mean, logvar)
        x_recon = self.decoder(z)
        return x_recon, mean, logvar

# 计算重构损失和KL散度
def loss_function(x_recon, x, mean, logvar):
    recon_loss = F.binary_cross_entropy(x_recon, x.view(-1, 784), reduction='sum')
    kl_loss = -0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp())
    return recon_loss + kl_loss

# 创建VAE模型实例
latent_dim = 20
vae = VAE(latent_dim)

# 定义优化器
optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)

# 训练模型
def train(model, optimizer, train_loader, num_epochs):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0
        for batch_idx, (x, _) in enumerate(train_loader):
            optimizer.zero_grad()
            x_recon, mean, logvar = model(x)
            loss = loss_function(x_recon, x, mean, logvar)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        avg_loss = total_loss / len(train_loader.dataset)
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, avg_loss))

# 使用MNIST数据集进行训练
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader

# 定义数据转换
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载MNIST数据集
train_dataset = MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

# 开始训练
num_epochs = 10
train(vae, optimizer, train_loader, num_epochs)

这是一个简单的VAE实现示例，使用了PyTorch中的`nn.Module`、`nn.Linear`和各种损失函数。编码器和解码器网络可以根据需求进行调整。训练过程中使用了MNIST数据集作为示例输入。在训练过程中，模型将逐渐学习生成近似输入图像的概率分布。

希望这个示例对你有所帮助！<span class="em">1</span><span class="em">2</span><span class="em">3</span>
#### 引用[.reference_title]
- *1* *2* *3* [【Pytorch深度学习实战】（11）变分自动编码器（VAE）](https://blog.csdn.net/sikh_0529/article/details/126939381)[target="_blank" data-report-click={"spm":"1018.2226.3001.9630","extra":{"utm_source":"vip_chatgpt_common_search_pc_result","utm_medium":"distribute.pc_search_result.none-task-cask-2~all~insert_cask~default-1-null.142^v92^chatsearchT0_1"}}] [.reference_item style="max-width: 100%"]
[ .reference_list ]