PyTorchで深層生成モデル#5: VAEの実装方法

【pytorchで深層生成モデル#5】pytorchでVAE

Variational Autoencoder(VAE)は、ニューラルネットワークを用いた深層生成モデルの一種であり、画像生成や次元削減などの様々な応用がある人気のあるモデルです。本チュートリアルでは、pytorchを使用してVAEを実装する方法について詳しく説明します。

1. VAEの原理

VAEは、生成モデルとして確率モデルを用い、潜在変数を導入してデータの分布を推定します。VAEの構造は、エンコーダーネットワークとデコーダーネットワークから成ります。エンコーダーネットワークは入力データを学習し、潜在変数の平均と分散を出力します。デコーダーネットワークは、潜在変数を入力として受け取り、生成したいデータの分布を出力します。

VAEの目的関数は、再構成誤差と正則化項で構成されており、以下のように表されます。

L = -ELBO = – E[logP(x|z)] + D_KL(Q(z|x) || P(z))

ここで、-ELBOはEvidence Lower Boundと呼ばれます。第1項は再構成誤差であり、生成したデータと元のデータの差を最小化することを目的とします。第2項はKLダイバージェンスと呼ばれ、潜在変数の分布Q(z|x)と事前分布P(z)の違いを最小化します。

1. データの前処理

まず、pytorchと必要なライブラリをインストールします。

pip install torch torchvision

次に、MNISTデータセットをダウンロードして前処理を行います。

import torch
from torchvision import datasets, transforms

# データの前処理
transform = transforms.Compose([
    transforms.ToTensor(),
])

# MNISTデータセットのダウンロード
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# データローダーの作成
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)

1. VAEの実装

VAEのエンコーダーとデコーダーをpytorchで実装します。以下は、簡単な例です。

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

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=512, latent_dim=2):
        super(VAE, self).__init__()

        # エンコーダー
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc21 = nn.Linear(hidden_dim, latent_dim)
        self.fc22 = nn.Linear(hidden_dim, latent_dim)

        # デコーダー
        self.fc3 = nn.Linear(latent_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, input_dim)

    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 mu + eps*std

    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

1. 学習

VAEの学習を行います。以下は、訓練とテストのループの例です。

model = VAE()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

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

def train(epoch):
    model.train()
    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()

def test(epoch):
    model.eval()
    test_loss = 0
    with torch.no_grad():
        for i, (data, _) in enumerate(test_loader):
            recon_batch, mu, logvar = model(data)
            test_loss += loss_function(recon_batch, data, mu, logvar).item()

for epoch in range(1, 51):
    train(epoch)
    test(epoch)

1. 結果の表示

学習が終わったら、VAEによって生成された画像を表示します。

import matplotlib.pyplot as plt
import numpy as np

# テストデータを使って画像生成
data, _ = next(iter(test_loader))
recon_batch, _, _ = model(data)
recon_batch = recon_batch.view(-1, 1, 28, 28)

# 元の画像と生成画像を表示
plt.figure(figsize=(10, 3))
for i in range(5):
    plt.subplot(2, 5, i+1)
    plt.imshow(data[i].numpy().reshape(28, 28), cmap='gray')
    plt.subplot(2, 5, i+6)
    plt.imshow(recon_batch[i].detach().numpy().reshape(28, 28), cmap='gray')
plt.show()

以上で、pytorchを使用してVAEを実装する方法についてのチュートリアルが終了しました。VAEはデーん生成や異常検知など、様々な応用があるため、ぜひ自分のプロジェクトに取り入れてみてください。