Полное руководство по коду с использованием Comet ML
Введение
В эпоху больших данных сжатие изображений стало необходимым для снижения затрат на хранение и передачу без ущерба для качества изображения. Традиционные методы сжатия, такие как JPEG, GIF и PNG, используют методы с потерями или без потерь, которые используют избыточность в значениях пикселей. Хотя эти методы эффективны, они ограничены тем, сколько данных они могут сжать.
Методы сжатия, основанные на глубоком обучении, стали жизнеспособной альтернативой традиционным методам. Одним из таких методов является Вариационный автоэнкодер (VAE), генеративная модель, которая изучает низкоразмерное представление изображения в скрытом пространстве. В этом руководстве рассматривается использование вариационных автоэнкодеров для сжатия изображений, принципы их работы и способы реализации модели VAE с использованием PyTorch. Мы также будем использовать Comet ML для отслеживания и регистрации производительности нашей модели вариационного автоэнкодера (VAE).
Что такое вариационные автоэнкодеры (VAE)?
Как правило, автоэнкодеры — это класс нейронных сетей, которые могут научиться сжимать и распаковывать данные с помощью итеративного процесса оптимизации. Они состоят из двух частей: кодировщика и декодера. Кодер отображает входные данные в представление более низкой размерности, в то время как декодер отображает представление более низкой размерности обратно в исходные данные. Автоэнкодеры можно обучать без присмотра, то есть им не требуются размеченные данные.
Вариационный автоэнкодер (VAE) — это тип автоэнкодера, который использует вероятностные методы для изучения сжатого представления входных данных. Основная идея заключается в том, что VAE использует кодировщик для преобразования входных данных высокой размерности в представление более низкой размерности и декодер для восстановления исходных данных из представления более низкой размерности.
Джозеф Рокка в статье под названием Понимание вариационных автоэнкодеров (VAE) определил вариационный автоэнкодер (VAE) как:
"автоэнкодер, обучение которого упорядочено, чтобы избежать переобучения и гарантировать, что скрытое пространство обладает хорошими свойствами, которые позволяют генеративный процесс".
В отличие от традиционных автоэнкодеров, которые кодируют ввод как единую точку, VAE изучают вероятностную модель распределения данных в скрытом пространстве. VAE моделируют распределение данных, используя распределение Гаусса, где кодировщик изучает среднее значение и дисперсию. Декодер генерирует выборку из изученного распределения путем выборки из распределения Гаусса с использованием среднего значения и дисперсии.
Рокка объясняет метод обучения модели в скрытом пространстве следующим образом:
- Вход кодируется как распределение по скрытому пространству.
- Из этого распределения выбирается точка из скрытого пространства.
- Точка выборки декодируется, и можно вычислить ошибку реконструкции.
- Наконец, ошибка реконструкции распространяется обратно по сети.
VAE изучают распределение вероятностей по низкоразмерному представлению вместо одного значения. Это позволяет им генерировать новые выборки данных путем выборки из изученного распределения вероятностей. VAE также оказались идеальными для нескольких задач, включая создание изображений и улучшение изображений с низким разрешением.
Анализ моделей в режиме реального времени позволяет вашей команде отслеживать, контролировать и корректировать модели, уже находящиеся в производстве. Узнайте больше уроков на местах с экспертами Comet.
Реализация VAE для сжатия изображений
При сжатии изображений мы можем использовать VAE для изучения сжатого представления изображения. Кодер отображает изображение в представление более низкого измерения, а декодер отображает представление более низкого измерения обратно в исходное изображение. Сжатое представление сохраняется или передается эффективно.
Давайте рассмотрим, как реализовать VAE для сжатия изображений с помощью PyTorch и набора данных MNIST.
Требования
Для этого урока вам понадобится следующее:
- Базовые знания Python и глубокое обучение.
- PyTorch и Comet ML: мы будем использовать их для реализации VAE. См. Документация по Comet, чтобы понять, как интегрировать Comet с PyTorch.
- Набор данных MNIST: набор данных, который содержит 60 000 небольших квадратных изображений 28x28 в градациях серого рукописных цифр от 0 до 9. Загрузить набор данных MNIST здесь
- Текстовые редакторы: подойдет Visual Studio Code или Sublime Text.
Выполнение
Установка и импорт библиотек
Мы начнем с настройки Comet ML для отслеживания экспериментов. Comet ML — платформа для отслеживания и воспроизводимости экспериментов в машинном обучении. Comet упрощает отслеживание и сравнение различных экспериментов, визуализацию результатов и совместную работу. Мы будем использовать Comet для обучения модели VAE в течение нескольких эпох и оценки ее производительности на проверочном наборе.
Во-первых, нам нужно создать учетную запись на Comet ML и установить пакет Comet ML:
!pip install comet_ml
Импортируйте необходимые библиотеки для проекта:
import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision.datasets import MNIST from torchvision.transforms import transforms from torch.autograd import Variable from comet_ml import Experiment
Перейдите к началу эксперимента с кометой:
from comet_ml import Experiment experiment = Experiment(api_key="your_api_key", project_name="vae-project")
Замените your_api_key своим ключом API Comet ML (это можно найти на странице настроек вашей учетной записи Comet ML). Замените vae-project на название вашего проекта.
Загрузка и предварительная обработка данных
Мы загрузим набор данных MNIST и предварительно обработаем данные, масштабируя их до диапазона [-1, 1].
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_data = MNIST(root='./data', train=True, transform=transform, download=True)
test_data = MNIST(root='./data', train=False, transform=transform, download=True)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
test_loader = DataLoader(test_data, batch_size=128, shuffle=False)
Определите модель
Давайте определим Encoder и Decoder как отдельные модули и объединим их, чтобы сформировать полную модель VAE. Кодер берет входное изображение и выводит среднее значение и дисперсию распределения Гаусса. Декодер берет образец из дистрибутива и генерирует реконструированное изображение.
Прямой метод модели VAE берет входное изображение и передает его через кодировщик, чтобы получить среднее значение и дисперсию распределения. Затем он производит выборку из распределения, используя прием репараметризации, и пропускает выборку через декодер для получения реконструированного изображения. Мы определим модель VAE с помощью PyTorch следующим образом:
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=1, padding=1)
self.fc1 = nn.Linear(7 * 7 * 64, 256)
self.fc21 = nn.Linear(256, 128)
self.fc22 = nn.Linear(256, 128)
def forward(self, x):
x = nn.functional.relu(self.conv1(x))
x = nn.functional.relu(self.conv2(x))
x = x.view(-1, 7 * 7 * 64)
x = nn.functional.relu(self.fc1(x))
mu = self.fc21(x)
logvar = self.fc22(x)
return mu, logvar
class Decoder(nn.Module):
def __init__(self):
super(Decoder, self).__init__()
self.fc1 = nn.Linear(128, 256)
self.fc2 = nn.Linear(256, 7 * 7 * 64)
self.conv1 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.ConvTranspose2d(32, 1, kernel_size=3, stride=1, padding=1)
def forward(self, z):
z = nn.functional.relu(self.fc1(z))
z = nn.functional.relu(self.fc2(z))
z = z.view(-1, 64, 7, 7)
z = nn.functional.relu(self.conv1(z))
z = nn.functional.tanh(self.conv2(z))
return z
class VAE(nn.Module):
def __init__(self):
super(VAE, self).__init__()
self.encoder = Encoder()
self.decoder = Decoder()
def forward(self, x):
mu, logvar = self.encoder(x)
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + eps * std
recon_x = self.decoder(z)
return recon_x, mu, logvar
Определите функцию потерь
Функция потерь состоит из двух терминов: потери binary cross-entropy (BCE), которая измеряет разницу между исходным изображением и реконструированным изображением, и расхождения Kullback-Leibler (KL), которая измеряет разницу между закодированным распределением и стандартным нормальным распределением.
Давайте определим функцию потерь ниже:
def vae_loss(recon_x, x, mu, logvar):
BCE = nn.functional.binary_cross_entropy_with_logits(recon_x.view(-1, 784), x.view(-1, 784), reduction='sum')
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD
Общая потеря представляет собой сумму потери BCE и потери KL-дивергенции. Потеря BCE гарантирует, что реконструированное изображение будет похоже на исходное изображение, в то время как потеря расхождения KL гарантирует, что распределение, полученное кодером, близко к стандартному нормальному распределению.
Инициализируйте модель и определите оптимизатор
Чтобы инициализировать модель VAE, мы создаем экземпляр объекта VAE с желаемыми гиперпараметрами, такими как размер скрытого пространства, количество слоев в кодере и декодере и количество каналов во входных и выходных изображениях. Мы также определяем оптимизатор, который будет использоваться для обновления параметров модели во время обучения. В этом примере мы используем оптимизатор Adam со скоростью обучения 0,001.
Вот код для инициализации модели и оптимизатора:
# Define hyperparameters latent_size = 20 input_channels = 1 output_channels = 1 hidden_dims = [32, 64, 128, 256] # Initialize model and move it to device model = VAE(latent_size, input_channels, output_channels, hidden_dims).to(device) # Define optimizer optimizer = optim.Adam(model.parameters(), lr=0.001)
Обучите модель
Давайте обучим модель, перебрав количество эпох и установив модель в режим обучения. Затем мы перебираем обучающие данные и перемещаем их на устройство. Мы должны установить градиенты равными нулю, передать входные данные через модель, чтобы получить реконструированное изображение, а также среднее значение и дисперсию распределения. Теперь мы вычисляем потери VAE и применяем обратное распространение градиентов. Кроме того, мы добавляем потери к общим потерям при обучении и обновляем веса модели с помощью оптимизатора.
for epoch in range(num_epochs):
model.train()
train_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
optimizer.zero_grad()
recon_batch, mu, logvar = model(data)
loss = vae_loss(recon_batch, data, mu, logvar)
loss.backward()
train_loss += loss.item()
optimizer.step()
if batch_idx % log_interval == 0:
current_loss = train_loss / (batch_idx + 1)
print('Epoch: {} [{}/{} ({:.0f}%)]\tTraining Loss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader),
current_loss))
experiment.log_metric('train_loss', current_loss, step=(epoch * len(train_loader) + batch_idx))
Теперь мы можем регистрировать текущие потери при обучении на консоли и в Comet, используя объект эксперимента. Кроме того, мы используем индекс партии и номер эпохи, чтобы отслеживать ход обучения.
Оценить модель
Пришло время оценить нашу обученную модель VAE на тестовом наборе данных MNIST и визуализировать реконструированные изображения. Для этого давайте установим модель в режим оценки и перейдем к тестовым данным. После этого переместите тестовые данные на устройство и пропустите их через модель, чтобы получить реконструированное изображение, а также среднее значение и дисперсию распределения. Затем мы вычисляем потери VAE и добавляем их к общим потерям при тестировании.
model.eval()
test_loss = 0
with torch.no_grad():
for data, _ in test_loader:
data = data.to(device)
recon_batch, mu, logvar = model(data)
test_loss += vae_loss(recon_batch, data, mu, logvar).item()
test_loss /= len(test_loader.dataset)
print('Test Loss: {:.6f}'.format(test_loss))
experiment.log_metric('test_loss', test_loss)
# visualize the reconstructed images
num_images = 8
for i in range(num_images):
fig, axs = plt.subplots(1, 2)
axs[0].imshow(data[i][0], cmap='gray')
axs[0].set_title('Original Image')
axs[1].imshow(recon_batch[i][0].cpu().detach().numpy(), cmap='gray')
axs[1].set_title('Reconstructed Image')
experiment.log_figure(figure_name=f"Reconstructed Image {i+1}", figure=plt)
Кроме того, мы вычисляем средние потери при тестировании и регистрируем их в консоли и Comet, используя объект эксперимента. Затем мы визуализируем несколько примеров исходных изображений и соответствующих им реконструированных изображений с помощью Matplotlib.
В этом примере мы создадим подграфик с двумя изображениями — одно для исходного изображения и одно для реконструированного изображения. Кроме того, давайте зарегистрируем каждое реконструированное изображение в Comet, используя метод log_figure объекта эксперимента.
Однако обратите внимание, что реконструированные изображения будут похожи на исходные изображения, но с некоторой потерей деталей из-за сжатия. Степень сжатия контролируется размером скрытого пространства, которое можно настроить, изменив параметр latent_size в классе VAE.
Вот пример реконструированных изображений, созданных моделью VAE:
Как мы видим, реконструированные изображения не идеальны, но они отражают основные черты исходных изображений. Вы можете увеличить степень сжатия, уменьшив размер скрытого пространства, но это может привести к дальнейшей потере деталей в восстановленных изображениях.
Заключение
В этом руководстве рассматриваются вариационные автоэнкодеры (VAE) для сжатия изображений. Мы реализовали модель VAE с помощью PyTorch и обучили ее набору данных MNIST. Мы использовали Comet для отслеживания и регистрации производительности модели во время обучения и оценки.
Мы показали, как VAE идеально подходят для сжатия изображений, отображая их в низкоразмерное скрытое пространство, а затем реконструируя их из этого пространства. VAE обеспечивают гибкую и мощную основу для сжатия изображений, которую можно адаптировать к широкому спектру наборов данных и приложений.
Я призываю вас поэкспериментировать с различными наборами данных, архитектурами и гиперпараметрами и изучить возможности и ограничения VAE для сжатия изображений.
Ресурсы
- Дидерик П. Кингма и Макс Веллинг (2019), Введение в вариационные автоэнкодеры || Основы и тенденции машинного обучения.
- Джозеф Рокка и Батист Рокка (2021), Понимание вариационных автоэнкодеров (VAE) || На пути к науке о данных.
- Комета Документы
- Набор данных MNIST
- Гауссово распределение || НаукаПрямой
Примечание редактора. Heartbeat — это интернет-издание и сообщество, созданное участниками и посвященное предоставлению лучших образовательных ресурсов для специалистов по науке о данных, машинному обучению и глубокому обучению. Мы стремимся поддерживать и вдохновлять разработчиков и инженеров из всех слоев общества.
Независимая от редакции, Heartbeat спонсируется и публикуется Comet, платформой MLOps, которая позволяет специалистам по данным и командам машинного обучения отслеживать, сравнивать, объяснять и оптимизировать свои эксперименты. Мы платим нашим авторам и не продаем рекламу.
Если вы хотите внести свой вклад, перейдите к нашему призыву к участию. Вы также можете подписаться на получение нашего еженедельного информационного бюллетеня (Еженедельник глубокого обучения), заглянуть в блог Comet, присоединиться к нам в Slack и подписаться на Comet в Twitter и LinkedIn для получения ресурсов, событий и многое другое, что поможет вам быстрее создавать лучшие модели машинного обучения.
