CNN Variational Autoencoder
PyTorch CNN VAEのサンプルコード解説
MNISTデータセットを使ったCNN版Variational Autoencoder（VAE）のコードを解説します。
完全なサンプルコード
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
# ハイパーパラメータ
BATCH_SIZE = 128
EPOCHS = 10
LEARNING_RATE = 1e-3
LATENT_DIM = 20 # 潜在変数の次元数
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# データセットの準備
transform = transforms.Compose([
transforms.ToTensor(),
])
train_dataset = datasets.MNIST('~/.pytorch/data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
# VAEモデルの定義
class VAE(nn.Module):
def __init__(self, latent_dim=20):
super(VAE, self).__init__()
# エンコーダー（画像 → 潜在変数）
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1), # 28x28 -> 14x14
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 14x14 -> 7x7
nn.ReLU(),
nn.Flatten(), # 64*7*7 = 3136
)
# 潜在変数の平均と分散を出力
self.fc_mu = nn.Linear(64*7*7, latent_dim)
self.fc_logvar = nn.Linear(64*7*7, latent_dim)
# デコーダー（潜在変数 → 画像）
self.decoder_input = nn.Linear(latent_dim, 64*7*7)
self.decoder = nn.Sequential(
nn.Unflatten(1, (64, 7, 7)), # 3136 -> 64x7x7
nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1), # 7x7 -> 14x14
nn.ReLU(),
nn.ConvTranspose2d(32, 1, kernel_size=3, stride=2, padding=1, output_padding=1), # 14x14 -> 28x28
nn.Sigmoid(), # [0, 1]の範囲に正規化
)
def encode(self, x):
"""エンコーダー部分"""
h = self.encoder(x)
mu = self.fc_mu(h)
logvar = self.fc_logvar(h)
return mu, logvar
def reparameterize(self, mu, logvar):
"""再パラメータ化トリック: z = μ + σ * ε"""
std = torch.exp(0.5 * logvar) # 標準偏差
eps = torch.randn_like(std) # 標準正規分布からサンプリング
z = mu + eps * std
return z
def decode(self, z):
"""デコーダー部分"""
h = self.decoder_input(z)
reconstruction = self.decoder(h)
return reconstruction
def forward(self, x):
"""順伝播"""
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
reconstruction = self.decode(z)
return reconstruction, mu, logvar
# 損失関数の定義
def vae_loss(recon_x, x, mu, logvar):
"""
VAEの損失関数 = 再構成誤差 + KLダイバージェンス
"""
# 再構成誤差（Binary Cross Entropy）
recon_loss = F.binary_cross_entropy(recon_x, x, reduction='sum')
# KLダイバージェンス
# KL(N(μ,σ²) || N(0,1)) = -0.5 * Σ(1 + log(σ²) - μ² - σ²)
kl_divergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return recon_loss + kl_divergence
# モデルの初期化
model = VAE(latent_dim=LATENT_DIM).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
# 学習ループ
def train(epoch):
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 % 100 == 0:
print(f'Epoch {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] '
f'Loss: {loss.item() / len(data):.4f}')
avg_loss = train_loss / len(train_loader.dataset)
print(f'====> Epoch: {epoch} Average loss: {avg_loss:.4f}')
# 学習実行
for epoch in range(1, EPOCHS + 1):
train(epoch)
# 生成画像の可視化
model.eval()
with torch.no_grad():
# ランダムサンプリングから生成
z = torch.randn(64, LATENT_DIM).to(device)
sample = model.decode(z).cpu()
# 画像表示
fig, axes = plt.subplots(8, 8, figsize=(10, 10))
for i, ax in enumerate(axes.flat):
ax.imshow(sample[i].squeeze(), cmap='gray')
ax.axis('off')
plt.tight_layout()
plt.savefig('vae_generated_samples.png')
plt.show()
# 再構成画像の可視化
with torch.no_grad():
data, _ = next(iter(train_loader))
data = data[:8].to(device)
recon, _, _ = model(data)
# 元画像と再構成画像を比較
fig, axes = plt.subplots(2, 8, figsize=(15, 4))
for i in range(8):
# 元画像
axes[0, i].imshow(data[i].cpu().squeeze(), cmap='gray')
axes[0, i].axis('off')
# 再構成画像
axes[1, i].imshow(recon[i].cpu().squeeze(), cmap='gray')
axes[1, i].axis('off')
axes[0, 0].set_ylabel('Original', size=20)
axes[1, 0].set_ylabel('Reconstructed', size=20)
plt.tight_layout()
plt.savefig('vae_reconstruction.png')
plt.show()
主要な構成要素の解説
1. エンコーダー（Encoder）
self.encoder = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
nn.ReLU(),
nn.Flatten(),
)
- 入力画像を畳み込み層で特徴抽出
- 28×28 → 14×14 → 7×7 と縮小
- 平均（μ）と分散（logvar）を出力
2. 再パラメータ化トリック
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
z = mu + eps * std
return z
- 目的: 確率的なサンプリングでも勾配が伝播できるようにする
- ε ~ N(0,1) を使って z = μ + σε と変換
3. デコーダー（Decoder）
self.decoder = nn.Sequential(
nn.ConvTranspose2d(64, 32, ...),
nn.ConvTranspose2d(32, 1, ...),
nn.Sigmoid(),
)
- 潜在変数から画像を再構成
- ConvTranspose2d（転置畳み込み）で画像を拡大
4. VAE損失関数
loss = 再構成誤差 + KLダイバージェンス
- 再構成誤差: 元画像と再構成画像の差
- KLダイバージェンス: 潜在変数分布をN(0,1)に近づける正則化項
実行結果
このコードを実行すると：