Le processus de mise en œuvre de la compression d'image : auto-encodeur variationnel

Variational Autoencoder (VAE) est un réseau neuronal d'apprentissage non supervisé utilisé pour la compression et la génération d'images. Par rapport aux encodeurs automatiques traditionnels, VAE peut reconstruire les images d’entrée et générer de nouvelles images similaires. L'idée principale est d'encoder l'image d'entrée dans une distribution de variables latentes et d'en échantillonner pour générer de nouvelles images. VAE est unique dans l'utilisation de l'inférence variationnelle pour entraîner le modèle, réalisant l'apprentissage des paramètres en maximisant la limite inférieure entre les données observées et générées. Cette méthode permet à VAE d’apprendre la structure sous-jacente des données et de générer de nouveaux échantillons. La VAE a obtenu un succès remarquable dans de nombreux domaines, notamment dans des tâches telles que la génération d'images, l'édition d'attributs et la reconstruction d'images.

VAE (autoencodeur variationnel) a une structure similaire à un autoencodeur et se compose d'un encodeur et d'un décodeur. Le codeur compresse l'image d'entrée en une distribution de variables latentes, comprenant un vecteur moyen et un vecteur de variance. Le décodeur échantillonne les variables latentes pour générer de nouvelles images. Afin de rendre la distribution des variables latentes plus raisonnable, VAE introduit le terme de régularisation de divergence KL pour rapprocher la distribution des variables latentes de la distribution normale standard. Cela peut améliorer l’expressivité et les capacités de génération du modèle.

Ce qui suit prend comme exemple l'ensemble de données de chiffres manuscrits du MNIST pour présenter le processus de mise en œuvre de la VAE.

Tout d'abord, nous devons importer les bibliothèques et les ensembles de données nécessaires.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable

# 加载数据集
transform = transforms.Compose([
    transforms.ToTensor(),
])
train_dataset = datasets.MNIST(root='./data/', train=True, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

Ensuite, définissez la structure du réseau de l'encodeur et du décodeur.

# 定义编码器
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        self.fc1 = nn.Linear(128 * 7 * 7, 256)
        self.fc21 = nn.Linear(256, 20) # 均值向量
        self.fc22 = nn.Linear(256, 20) # 方差向量

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.relu(self.conv3(x))
        x = x.view(-1, 128 * 7 * 7)
        x = nn.functional.relu(self.fc1(x))
        mean = self.fc21(x)
        log_var = self.fc22(x)
        return mean, log_var


# 定义解码器
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(20, 256)
        self.fc2 = nn.Linear(256, 128 * 7 * 7)
        self.conv1 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.conv2 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.conv3 = nn.ConvTranspose2d(32, 1, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        x = x.view(-1, 128, 7, 7)
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.sigmoid(self.conv3(x))
        return x


# 定义VAE模型
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()

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

    def forward(self, x):
        mean, log_var = self.encoder(x)

La prochaine étape est le processus de propagation vers l'avant du modèle VAE, qui comprend l'échantillonnage à partir de variables latentes pour générer de nouvelles images et le calcul des termes de régularisation de l'erreur de reconstruction et de la divergence KL.

z = self.reparameterize(mean, log_var)
x_recon = self.decoder(z)
return x_recon, mean, log_var

def loss_function(self, x_recon, x, mean, log_var):
    recon_loss = nn.functional.binary_cross_entropy(x_recon, x, size_average=False)
    kl_loss = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())
    return recon_loss + kl_loss

def sample(self, num_samples):
    z = torch.randn(num_samples, 20)
    samples = self.decoder(z)
    return samples

Enfin, nous définissons l'optimiseur et commençons à entraîner le modèle.

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

# 开始训练模型
num_epochs = 10
for epoch in range(num_epochs):
for batch_idx, (data, _) in enumerate(train_loader):
data = Variable(data)
optimizer.zero_grad()
x_recon, mean, log_var = vae(data)
loss = vae.loss_function(x_recon, data, mean, log_var)
loss.backward()
optimizer.step()

    if batch_idx % 100 == 0:
        print('Epoch [{}/{}], Batch [{}/{}], Loss: {:.4f}'.format(
            epoch+1, num_epochs, batch_idx+1, len(train_loader), loss.data.item()))

Une fois la formation terminée, nous pouvons utiliser la VAE pour générer de nouvelles images de chiffres manuscrits.

# 生成手写数字图像
samples = vae.sample(10)
fig, ax = plt.subplots(1, 10, figsize=(10, 1))
for i in range(10):
ax[i].imshow(samples[i].detach().numpy().reshape(28, 28), cmap='gray')
ax[i].axis('off')
plt.show()

VAE est un puissant modèle de compression d'image et génératif qui réalise la compression d'image en codant les images d'entrée dans une distribution de variables latentes tout en les échantillonnant pour générer de nouvelles images. Différent des auto-encodeurs traditionnels, VAE introduit également le terme de régularisation de divergence KL pour rendre la distribution des variables latentes plus raisonnable. Lors de la mise en œuvre de VAE, il est nécessaire de définir la structure du réseau du codeur et du décodeur, et de calculer les termes de régularisation de l'erreur de reconstruction et de la divergence KL. En entraînant le modèle VAE, la distribution des variables latentes de l'image d'entrée peut être apprise et de nouvelles images peuvent être générées à partir de celle-ci.

Ce qui précède est le processus de base d'introduction et de mise en œuvre de la VAE. J'espère qu'il sera utile aux lecteurs.

Ce qui précède est le contenu détaillé de. pour plus d'informations, suivez d'autres articles connexes sur le site Web de PHP en chinois!