VAE自编码器案例

蔚落 2023-09-30 21:57 11阅读 0赞

#### VAE自编码器案例 ####

下面是一个使用PyTorch实现变分自编码器（VAE）的案例。VAE是一种更复杂的自编码器，它可以生成新的样本，并且在对潜在空间进行插值时呈现出更好的连续性。

首先导入必要的库和模块：

import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms
    import torch.nn.functional as F

接着定义一个类来表示我们的VAE：

class VAE(nn.Module):
        def __init__(self, latent_dim=20):
            super(VAE, self).__init__()
    
            # Encoder
            self.fc1 = nn.Linear(784, 512)
            self.fc21 = nn.Linear(512, latent_dim)  # mean
            self.fc22 = nn.Linear(512, latent_dim)  # variance
    
            # Decoder
            self.fc3 = nn.Linear(latent_dim, 512)
            self.fc4 = nn.Linear(512, 784)
    
        def encode(self, x):
            h1 = F.relu(self.fc1(x))
            return self.fc21(h1), self.fc22(h1)
    
        def reparameterize(self, mu, logvar):
            std = torch.exp(0.5*logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
    
        def decode(self, z):
            h3 = F.relu(self.fc3(z))
            return torch.sigmoid(self.fc4(h3))
    
        def forward(self, x):
            mu, logvar = self.encode(x.view(-1, 784))
            z = self.reparameterize(mu, logvar)
            return self.decode(z), mu, logvar

这个VAE包含了一个编码器和一个解码器，其具体实现细节如下：

编码器通过两个全连接层将输入图像压缩为隐变量的均值和方差。  
解码器通过两个全连接层，并应用sigmoid函数将隐变量重构为原始图像。  
关键方法包括：

encode(x)：输入x，输出其均值和方差；
    reparameterize(mu, logvar)：通过重新参数化技巧建立了梯度传递的路径；
    decode(z)：输入隐变量z，返回重构后的原始图像。

接下来定义一些超参数：

batch_size = 128
    learning_rate = 1e-3
    num_epochs = 10
    latent_dim = 20

然后加载MNIST数据集并进行预处理：

transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    
    train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

在这里，我们使用了transforms库来定义一个变换，它将图像转换为张量，并归一化其像素值以使其范围在 -1 到 1 之间。我们还定义了一个DataLoader对象，用于批量读取数据。

接下来初始化模型和优化器：

model = VAE(latent_dim=latent_dim)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

在这里，我们使用了Adam优化器作为优化算法。

接下来定义损失函数：

def loss_function(recon_x, x, mu, logvar):
        BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return BCE + KLD

在这里，我们定义了一个损失函数，它由两部分组成。第一项是重构误差，通过计算输入图像x的二进制交叉熵损失函数得到。第二项是KL散度，用于惩罚潜在变量z不符合标准正态分布的程度。

最后训练完模型后，我们可以使用生成方法来检查VAE的效果。以下代码示例展示了如何生成具有潜在变量z的新图像：

继续训练模型：

for epoch in range(num_epochs):
        train_loss = 0
        for batch_idx, (data, _) in enumerate(train_loader):
            optimizer.zero_grad()
            recon_batch, mu, logvar = model(data)
            loss = loss_function(recon_batch, data, mu, logvar)
            loss.backward()
            train_loss += loss.item()
            optimizer.step()
    
        print('Epoch: {}  Loss: {:.3f}'.format(epoch+1, train_loss/len(train_loader.dataset)))

在这里，我们对于每个epoch，遍历训练数据集，并将每个batch的数据加载到模型中。首先将图像展平为一维张量，通过VAE进行前向传播得到重构后的图像和隐变量的均值和方差，然后计算重构误差和KL散度并执行反向传播更新权重。

最后打印出当前epoch的平均损失值。

除了训练，我们还可以使用生成方法来检查VAE的效果。训练完模型后，我们可以使用生成方法来检查VAE的效果。以下代码示例展示了如何生成具有潜在变量z的新图像：

# 随机采样10个隐变量并解码
    
    
    with torch.no_grad():
        z = torch.randn(10, latent_dim).to(device)
        sample = model.decode(z).cpu()
        save_image(sample.view(10, 1, 28, 28), './samples/sample_' + str(epoch) + '.png')

在这里，我们首先使用torch.randn()函数从标准正态分布中采样隐变量。然后，我们将其传递给解码器，生成10个新的图像并保存为PNG文件。

这就是一个使用PyTorch实现的简单变分自编码器（VAE）示例。

完整的训练代码如下：

import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader
    from torchvision import datasets, transforms
    import torch.nn.functional as F
    from torchvision.utils import save_image
    
    class VAE(nn.Module):
        def __init__(self, latent_dim=20):
            super(VAE, self).__init__()
    
            # Encoder
            self.fc1 = nn.Linear(784, 512)
            self.fc21 = nn.Linear(512, latent_dim)  # mean
            self.fc22 = nn.Linear(512, latent_dim)  # variance
    
            # Decoder
            self.fc3 = nn.Linear(latent_dim, 512)
            self.fc4 = nn.Linear(512, 784)
    
        def encode(self, x):
            h1 = F.relu(self.fc1(x))
            return self.fc21(h1), self.fc22(h1)
    
        def reparameterize(self, mu, logvar):
            std = torch.exp(0.5*logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
    
        def decode(self, z):
            h3 = F.relu(self.fc3(z))
            return torch.sigmoid(self.fc4(h3))
    
        def forward(self, x):
            mu, logvar = self.encode(x.view(-1, 784))
            z = self.reparameterize(mu, logvar)
            return self.decode(z), mu, logvar
    
    batch_size = 128
    learning_rate = 1e-3
    num_epochs = 10
    latent_dim = 20
    
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    
    train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = VAE(latent_dim=latent_dim).to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    def loss_function(recon_x, x, mu, logvar):
        BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        return BCE + KLD
    
    for epoch in range(num_epochs):
        train_loss = 0
        for batch_idx, (data, _) in enumerate(train_loader):
            optimizer.zero_grad()
            recon_batch, mu, logvar = model(data.to(device))
            loss = loss_function(recon_batch, data.to(device), mu, logvar)
            loss.backward()
            train_loss += loss.item()
            optimizer.step()
    
        print('Epoch: {}  Loss: {:.3f}'.format(epoch+1, train_loss/len(train_loader.dataset)))
    
        # 随机采样10个隐变量并解码
        with torch.no_grad():
            z = torch.randn(10, latent_dim).to(device)
            sample = model.decode(z).cpu()
            save_image(sample.view(10, 1, 28, 28), './samples/sample_' + str(epoch) + '.png')

在训练过程中，我们打印出每个epoch的平均损失值，并在每个epoch结束时生成10个新图像。这些图像被保存在./samples文件夹中。