В этой статье мы будем реализовывать вариационные автоэнкодеры с нуля на python.

Что такое автоэнкодеры и для чего они служат

Автоэнкодер — это нейронная архитектура, состоящая из двух частей: кодера и декодера. Декодер следует за кодировщиком, а посередине находится так называемый скрытый слой, который имеет различные названия, иногда может упоминаться как слой узкого места, скрытое пространство, скрытый слой, слой кодирования или слой кода. Это выглядит примерно так:

Архитектура автоэнкодера

Автоэнкодеры служат различным целям. Наверное, самое очевидное — это сжатие данных: очевидно, что при прохождении входного сигнала через кодирующую часть скрытое представление изображения значительно меньше по размеру. Например, на изображении выше, хотя входной сигнал представлен 8 значениями, его сжатое представление занимает только 3.

Автоэнкодеры также могут служить различным другим целям: шумоподавление данных, изучение признаков, обнаружение аномалий и, как мы позже увидим в статье, генеративное моделирование. Вот несколько статей, которые могут быть вам интересны для прочтения:

  1. Хинтон, Г.Э., и Салахутдинов, Р.Р. (2006). Уменьшение размерности данных с помощью нейронных сетей. Наука, 313 (5786), 504–507. doi: 10.1126/science.1127647. Ссылка: http://www.cs.toronto.edu/~hinton/absps/science.pdf
  2. Бенжио, Ю., Ламблин, П., Поповичи, Д., и Ларошель, Х. (2007). Жадное послойное обучение глубоких сетей. Достижения в области систем обработки нейронной информации 19 (стр. 153–160). Ссылка: https://papers.nips.cc/paper/3048-greedy-layer-wise-training-of-deep-networks.pdf
  3. Ле, К.В., Нгиам, Дж., и Коутс, А. (2013). Создание функций высокого уровня с использованием крупномасштабного неконтролируемого обучения. На Международной конференции по машинному обучению (стр. 91–99). Ссылка: http://proceedings.mlr.press/v28/le13.pdf
  4. Кингма, Д. П., и Веллинг, М. (2014). Автокодирование вариационного Байеса. На Международной конференции по обучающим представлениям. Ссылка: https://arxiv.org/abs/1312.6114

Реализация автоэнкодеров

Мы будем работать с хорошо известным набором данных MNIST. Чтобы загрузить MNIST в локальную папку, выполните следующее:

# Download the files
url = "http://yann.lecun.com/exdb/mnist/"
filenames = ['train-images-idx3-ubyte.gz', 'train-labels-idx1-ubyte.gz',
             't10k-images-idx3-ubyte.gz', 't10k-labels-idx1-ubyte.gz']
data = []
for filename in filenames:
    print("Downloading", filename)
    request.urlretrieve(url + filename, filename)
    with gzip.open(filename, 'rb') as f:
        if 'labels' in filename:
            # Load the labels as a one-dimensional array of integers
            data.append(np.frombuffer(f.read(), np.uint8, offset=8))
        else:
            # Load the images as a two-dimensional array of pixels
            data.append(np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1,28*28))

# Split into training and testing sets
X_train, y_train, X_test, y_test = data

# Normalize the pixel values
X_train = X_train.astype(np.float32) / 255.0
X_test = X_test.astype(np.float32) / 255.0

# Convert labels to integers
y_train = y_train.astype(np.int64)
y_test = y_test.astype(np.int64)

Итак, теперь у нас есть набор для обучения и тестирования, и мы готовы реализовать архитектуру. Но сначала давайте просто посмотрим на наши изображения с этим:

def show_images(images, labels):
    """
    Display a set of images and their labels using matplotlib.
    The first column of `images` should contain the image indices,
    and the second column should contain the flattened image pixels
    reshaped into 28x28 arrays.
    """
    # Extract the image indices and reshaped pixels
    pixels = images.reshape(-1, 28, 28)

    # Create a figure with subplots for each image
    fig, axs = plt.subplots(
        ncols=len(images), nrows=1, figsize=(10, 3 * len(images))
    )

    # Loop over the images and display them with their labels
    for i in range(len(images)):
        # Display the image and its label
        axs[i].imshow(pixels[i], cmap="gray")
        axs[i].set_title("Label: {}".format(labels[i]))

        # Remove the tick marks and axis labels
        axs[i].set_xticks([])
        axs[i].set_yticks([])
        axs[i].set_xlabel("Index: {}".format(i))

    # Adjust the spacing between subplots
    fig.subplots_adjust(hspace=0.5)

    # Show the figure
    plt.show()

Архитектура автоэнкодера очень проста

import torch.nn as nn

class AutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Set the number of hidden units
        self.num_hidden = 8
        
        # Define the encoder part of the autoencoder
        self.encoder = nn.Sequential(
            nn.Linear(784, 256),  # input size: 784, output size: 256
            nn.ReLU(),  # apply the ReLU activation function
            nn.Linear(256, self.num_hidden),  # input size: 256, output size: num_hidden
            nn.ReLU(),  # apply the ReLU activation function
        )
        
        # Define the decoder part of the autoencoder
        self.decoder = nn.Sequential(
            nn.Linear(self.num_hidden, 256),  # input size: num_hidden, output size: 256
            nn.ReLU(),  # apply the ReLU activation function
            nn.Linear(256, 784),  # input size: 256, output size: 784
            nn.Sigmoid(),  # apply the sigmoid activation function to compress the output to a range of (0, 1)
        )

    def forward(self, x):
        # Pass the input through the encoder
        encoded = self.encoder(x)
        # Pass the encoded representation through the decoder
        decoded = self.decoder(encoded)
        # Return both the encoded representation and the reconstructed output
        return encoded, decoded

Чтобы обучить эту штуку, нам на самом деле не нужны метки изображений, потому что мы имеем дело с неконтролируемым подходом. Мы используем простую среднеквадратичную ошибку потери, так как мы хотим реконструировать наши изображения наиболее точным способом. Проведем подготовительную работу:

# Convert the training data to PyTorch tensors
X_train = torch.from_numpy(X_train)

# Create the autoencoder model and optimizer
model = AutoEncoder()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Define the loss function
criterion = nn.MSELoss()

# Set the device to GPU if available, otherwise use CPU
model.to(device)

# Create a DataLoader to handle batching of the training data
train_loader = torch.utils.data.DataLoader(
    X_train, batch_size=batch_size, shuffle=True
)

Наконец, цикл обучения выглядит очень просто и ничего неожиданного:

# Training loop
for epoch in range(num_epochs):
    total_loss = 0.0
    for batch_idx, data in enumerate(train_loader):
        # Get a batch of training data and move it to the device
        data = data.to(device)

        # Forward pass
        encoded, decoded = model(data)

        # Compute the loss and perform backpropagation
        loss = criterion(decoded, data)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Update the running loss
        total_loss += loss.item() * data.size(0)

    # Print the epoch loss
    epoch_loss = total_loss / len(train_loader.dataset)
    print(
        "Epoch {}/{}: loss={:.4f}".format(epoch + 1, num_epochs, epoch_loss)
    )

Чтобы рассчитать потери, мы просто берем входные изображения и сравниваем их с реконструированной версией. Обучение должно занять всего несколько секунд, и если мы сравним выходные и входные изображения, то увидим следующее:

Верхний ряд - это исходные изображения с реконструированными изображениями, изображенными в нижнем ряду.

Есть несколько интересных замечаний:

  1. Восстановленные изображения размыты. Это потому, что реконструкция не идеальна.
  2. Последняя колонка показывает, что этот тип сжатия платный, за него приходится платить ошибкой при декодировании.
  3. Мы используем только 8 скрытых единиц, увеличение количества скрытых единиц приведет к улучшению качества изображения, тогда как их уменьшение ухудшит размытие.

Вот. С 2 скрытыми элементами вместо 8:

Мало того, что мы видим значительно более размытые изображения, так еще и автоэнкодер делает некоторые очевидные ошибки, когда пытается декодировать скрытое пространство.

С 32 элементами в середине:

Интересно, что мы не видим больших улучшений по сравнению с использованием всего 8 нейронов посередине. Что не очень удивительно, хотя и трудно предсказать. В таких системах только эксперименты могут показать, что лучше. Похоже, что для этой конкретной задачи использование 32 единиц для скрытого пространства является излишним.

С концепцией выше все в порядке. Однако есть кое-что, в чем архитектура очень слаба. В частности, очень сложно генерировать новые данные. Подумайте об этом, если мы отбросим часть кодировщика и начнем только со скрытого слоя, мы сможем прийти к осмысленному изображению в конце. Однако неясно, как осмысленно сэмплировать скрытое пространство, то есть придумать надежную стратегию сэмплирования, которая гарантировала бы, что выходные изображения будут читаемыми, а не просто случайным шумом.

Хотя на самом деле можно сэмплировать 2- или 3-мерное пространство, если мы попытаемся работать с более сложными объектами, это может быть намного сложнее. Давайте попробуем это сделать. Прежде всего, вот все изображения в тренировочном наборе.

Теперь что мы делаем, так это генерируем кучу семплов из этого скрытого пространства. Немного сложно генерировать выборки из этого точного распределения скрытого пространства, поэтому вместо этого мы генерируем выборки из прямоугольника. Вот что мы получаем:

Выборка изученного скрытого пространства автоэнкодера.

В то время как некоторые образцы выглядят вполне нормально, будет труднее осмысленно сэмплировать одно и то же пространство, поскольку размерность этого пространства была выше. Например, если мы увеличим размерность до 32, вот что мы получим из подобного эксперимента:

Цифры больше не распознаются, и вам нужно намного дольше сэмплировать, чтобы получить что-то, что имеет визуальный смысл. Есть ли способ лучше?

Вариационные автоэнкодеры

Оригинальная статья, посвященная вариационным автоэнкодерам (VAE), называется Auto-Encoding Variational Bayes и была опубликована Дидериком П. Кингмой и Максом Веллингом в 2014 году. Вы можете найти статью на сервере препринтов arXiv по следующей ссылке: https: //arxiv.org/abs/1312.6114

VAE дают нам более гибкий способ изучения скрытого представления предметной области. Основная идея очень проста: мы узнаем параметры распределения скрытого пространства, из которого будем производить выборку. Вот как это выглядит на схеме:

Мы используем входные данные и изучаем векторы средних и дисперсий, чтобы позже использовать их для выборки скрытой переменной, которая должна быть декодирована с помощью декодера.

Это здорово, но есть проблема с настройкой, особенно с блоком, который сэмплирует скрытые переменные. Дело в том, что операция выборки является недифференцируемой. Чтобы избежать этой проблемы, существует так называемый прием перепараметризации. Это работает следующим образом: вместо того, чтобы получать выборки из этого блока, мы явно изучаем два вектора средних и дисперсий, а затем имеем независимый блок, который просто выдает выборки из

Поэтому вместо прямой выборки этих дистрибутивов мы делаем следующее:

Где 𝐿𝑖 представляют компонент скрытого представления. Итак, наша диаграмма теперь выглядит так:

Теперь наш процесс дифференцируем. Давайте закодируем это:

class VAE(AutoEncoder):
    def __init__(self):
        super().__init__()
        # Add mu and log_var layers for reparameterization
        self.mu = nn.Linear(self.num_hidden, self.num_hidden)
        self.log_var = nn.Linear(self.num_hidden, self.num_hidden)

    def reparameterize(self, mu, log_var):
        # Compute the standard deviation from the log variance
        std = torch.exp(0.5 * log_var)
        # Generate random noise using the same shape as std
        eps = torch.randn_like(std)
        # Return the reparameterized sample
        return mu + eps * std

    def forward(self, x):
        # Pass the input through the encoder
        encoded = self.encoder(x)
        # Compute the mean and log variance vectors
        mu = self.mu(encoded)
        log_var = self.log_var(encoded)
        # Reparameterize the latent variable
        z = self.reparameterize(mu, log_var)
        # Pass the latent variable through the decoder
        decoded = self.decoder(z)
        # Return the encoded output, decoded output, mean, and log variance
        return encoded, decoded, mu, log_var

    def sample(self, num_samples):
        with torch.no_grad():
            # Generate random noise
            z = torch.randn(num_samples, self.num_hidden).to(device)
            # Pass the noise through the decoder to generate samples
            samples = self.decoder(z)
        # Return the generated samples
        return samples

Теперь главное, как мы тренируем этот материал. Давайте определим нашу функцию потерь:

# Define a loss function that combines binary cross-entropy and Kullback-Leibler divergence
def loss_function(recon_x, x, mu, logvar):
    # Compute the binary cross-entropy loss between the reconstructed output and the input data
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction="sum")
    # Compute the Kullback-Leibler divergence between the learned latent variable distribution and a standard Gaussian distribution
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    # Combine the two losses by adding them together and return the result
    return BCE + KLD

Первая составляющая нам уже знакома — это просто ошибка реконструкции. Второй компонент вводит штраф за то, что изученное распределение слишком сильно отклоняется от предыдущего распределения. Наше априорное распределение — это всего лишь набор нормальных распределений единичной дисперсии с нулевым центром.

Итак, для обучения модели мы можем просто использовать:

def train_vae(X_train, learning_rate=1e-3, num_epochs=10, batch_size=32):
    # Convert the training data to PyTorch tensors
    X_train = torch.from_numpy(X_train).to(device)

    # Create the autoencoder model and optimizer
    model = VAE()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Define the loss function
    criterion = nn.MSELoss(reduction="sum")

    # Set the device to GPU if available, otherwise use CPU
    model.to(device)

    # Create a DataLoader to handle batching of the training data
    train_loader = torch.utils.data.DataLoader(
        X_train, batch_size=batch_size, shuffle=True
    )

    # Training loop
    for epoch in range(num_epochs):
        total_loss = 0.0
        for batch_idx, data in enumerate(train_loader):
            # Get a batch of training data and move it to the device
            data = data.to(device)

            # Forward pass
            encoded, decoded, mu, log_var = model(data)

            # Compute the loss and perform backpropagation
            KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
            loss = criterion(decoded, data) + 3 * KLD
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # Update the running loss
            total_loss += loss.item() * data.size(0)

        # Print the epoch loss
        epoch_loss = total_loss / len(train_loader.dataset)
        print(
            "Epoch {}/{}: loss={:.4f}".format(epoch + 1, num_epochs, epoch_loss)
        )

    # Return the trained model
    return model

Наконец, мы можем сгенерировать несколько изображений:

Обратите внимание, что изображения выглядят размытыми. Это происходит потому, что мы используем MAE для контроля качества реконструкции.

И последнее, что довольно интересно сделать, — это сгенерировать случайный вектор, а затем постепенно изменить один из его компонентов, сохраняя при этом другие компоненты фиксированными. Таким образом, мы можем видеть, как меняется вывод декодера. Вот несколько примеров:

Не обращайте внимания на подписи к ярлыкам. Обратите внимание, как мы сгенерировали вектор, затем, изменив один из его компонентов, мы перешли от 0 к 9, а затем к 1 или 7.

Заключение

В заключение, внедрение автоэнкодеров с нуля — полезное упражнение для понимания основ обучения без учителя и сжатия данных. В то время как простые автоэнкодеры могут реконструировать изображения, им сложно генерировать новые данные. Вариационные автоэнкодеры (VAE) предлагают более гибкий подход, изучая параметры распределения скрытого пространства, которые могут быть выбраны для генерации новых данных. Трюк перепараметризации используется, чтобы сделать операцию выборки дифференцируемой. VAE предлагают улучшение по сравнению с простыми автоэнкодерами и имеют множество приложений, таких как генеративное моделирование, обнаружение аномалий и изучение признаков.