Postsもっと読む
CNNによる二値画像AutoEncoder
PyTorch CNNによる二値画像AutoEncoderの解説
二値画像(白黒画像)を扱うAutoEncoderのサンプルコードを解説します。
完全なサンプルコード
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
# デバイスの設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# MPSが利用可能かチェック
if torch.backends.mps.is_available():
device = torch.device("mps")
# ハイパーパラメータ
BATCH_SIZE = 128
LEARNING_RATE = 0.001
EPOCHS = 3
LATENT_DIM = 32 # 潜在空間の次元数
# データの準備(MNISTを例に使用)
transform = transforms.Compose([
transforms.ToTensor(),
#transforms.Normalize((0.5,), (0.5,)) # -1~1の範囲に正規化
transforms.Lambda(lambda x: (x > 0.5).float()) # 二値化
])
train_dataset = datasets.MNIST(
root='~/.pytorch/data',
train=True,
download=True,
transform=transform
)
train_loader = DataLoader(
train_dataset,
batch_size=BATCH_SIZE,
shuffle=True
)
# Encoderの定義
class Encoder(nn.Module):
def __init__(self, latent_dim):
super(Encoder, self).__init__()
# 畳み込み層
self.conv_layers = nn.Sequential(
# 入力: 1x28x28
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1), # 32x14x14
nn.BatchNorm2d(32),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 64x7x7
nn.BatchNorm2d(64),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 128x4x4
nn.BatchNorm2d(128),
nn.ReLU(),
)
# 全結合層で潜在空間へ
self.fc = nn.Linear(128 * 4 * 4, latent_dim)
def forward(self, x):
x = self.conv_layers(x)
x = x.view(x.size(0), -1) # Flatten
x = self.fc(x)
return x
# Decoderの定義
class Decoder(nn.Module):
def __init__(self, latent_dim):
super(Decoder, self).__init__()
# 潜在空間から特徴マップへ
self.fc = nn.Linear(latent_dim, 128 * 4 * 4)
# 転置畳み込み層(逆畳み込み)
self.deconv_layers = nn.Sequential(
# 入力: 128x4x4
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # 64x8x8
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1), # 32x16x16
nn.BatchNorm2d(32),
nn.ReLU(),
nn.ConvTranspose2d(32, 1, kernel_size=4, stride=2, padding=3), # 1x28x28
#nn.Tanh() # -1~1の範囲に出力
nn.Sigmoid() # 0-1の範囲に出力を制限
)
def forward(self, x):
x = self.fc(x)
x = x.view(x.size(0), 128, 4, 4) # Reshape
x = self.deconv_layers(x)
return x
# AutoEncoderの定義
class AutoEncoder(nn.Module):
def __init__(self, latent_dim):
super(AutoEncoder, self).__init__()
self.encoder = Encoder(latent_dim)
self.decoder = Decoder(latent_dim)
def forward(self, x):
latent = self.encoder(x)
reconstructed = self.decoder(latent)
return reconstructed
# モデルの初期化
model = AutoEncoder(LATENT_DIM).to(device)
# 損失関数と最適化手法
# criterion = nn.MSELoss() # 平均二乗誤差
# 損失関数(二値交差エントロピー)
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
# 訓練ループ
def train(model, train_loader, criterion, optimizer, epochs):
model.train()
losses = []
for epoch in range(epochs):
epoch_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
data = data.to(device)
# 勾配の初期化
optimizer.zero_grad()
# 順伝播
reconstructed = model(data)
# 損失計算
loss = criterion(reconstructed, data)
# 逆伝播
loss.backward()
# パラメータ更新
optimizer.step()
epoch_loss += loss.item()
if batch_idx % 100 == 0:
print(f'Epoch [{epoch+1}/{epochs}], Step [{batch_idx}/{len(train_loader)}], Loss: {loss.item():.4f}')
avg_loss = epoch_loss / len(train_loader)
losses.append(avg_loss)
print(f'Epoch [{epoch+1}/{epochs}], Average Loss: {avg_loss:.4f}')
return losses
# 訓練実行
losses = train(model, train_loader, criterion, optimizer, EPOCHS)
# 結果の可視化
def visualize_results(model, test_loader, num_images=10):
model.eval()
with torch.no_grad():
data, _ = next(iter(test_loader))
data = data[:num_images].to(device)
reconstructed = model(data)
# CPU に移動して表示
data = data.cpu()
reconstructed = reconstructed.cpu()
fig, axes = plt.subplots(2, num_images, figsize=(15, 3))
for i in range(num_images):
# 元画像
axes[0, i].imshow(data[i].squeeze(), cmap='gray')
axes[0, i].axis('off')
if i == 0:
axes[0, i].set_title('Original', fontsize=10)
# 再構成画像
axes[1, i].imshow(reconstructed[i].squeeze(), cmap='gray')
axes[1, i].axis('off')
if i == 0:
axes[1, i].set_title('Reconstructed', fontsize=10)
plt.tight_layout()
plt.savefig('autoencoder_results.png')
plt.show()
# テストデータで可視化
test_dataset = datasets.MNIST(
root='~/.pytorch/data',
train=False,
download=True,
transform=transform
)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
visualize_results(model, test_loader)
# 損失の推移をプロット
plt.figure(figsize=(10, 5))
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training Loss')
plt.grid(True)
plt.savefig('autoencoder_loss.png')
plt.show()
主要部分の解説
1. Encoder(符号化器)
nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1)
- 畳み込み層で画像の特徴を抽出
stride=2で画像サイズを半分に縮小- 徐々にチャネル数を増やして特徴を豊かに
2. Decoder(復号化器)
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)
- 転置畳み込みで画像を拡大
- Encoderの逆操作を実行
- 最終的に元の画像サイズに復元
3. 損失関数
criterion = nn.MSELoss()
- 元画像と再構成画像の差を最小化
- 二値画像には
BCELossも使用可能
このコードを実行すると、画像の圧縮・復元が学習され、ノイズ除去や特徴抽出にも応用できます。
Postsもっと読む
PyTorchによる二値画像AutoEncoder
PyTorchによる二値画像AutoEncoderのサンプルコード解説
二値画像(白黒画像)を扱うAutoEncoderの実装例を解説します。
完全なサンプルコード
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
# 1. AutoEncoderモデルの定義
class BinaryAutoEncoder(nn.Module):
def __init__(self, input_dim=784, hidden_dim=128, latent_dim=32):
super(BinaryAutoEncoder, self).__init__()
# エンコーダ部分
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, latent_dim),
nn.ReLU()
)
# デコーダ部分
self.decoder = nn.Sequential(
nn.Linear(latent_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid() # 0-1の範囲に出力を制限
)
def forward(self, x):
# エンコード
encoded = self.encoder(x)
# デコード
decoded = self.decoder(encoded)
return decoded
def encode(self, x):
return self.encoder(x)
# 2. データの準備
def prepare_data(batch_size=128):
# MNISTデータセットを使用(二値化)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: (x > 0.5).float()) # 二値化
])
train_dataset = datasets.MNIST(
root='~/.pytorch/data',
train=True,
download=True,
transform=transform
)
test_dataset = datasets.MNIST(
root='~/.pytorch/data',
train=False,
download=True,
transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
return train_loader, test_loader
# 3. 訓練関数
def train(model, train_loader, optimizer, criterion, device):
model.train()
total_loss = 0
for batch_idx, (data, _) in enumerate(train_loader):
# データを平坦化 (batch_size, 1, 28, 28) -> (batch_size, 784)
data = data.view(data.size(0), -1).to(device)
# 勾配をゼロに
optimizer.zero_grad()
# 順伝播
output = model(data)
# 損失計算
loss = criterion(output, data)
# 逆伝播
loss.backward()
# パラメータ更新
optimizer.step()
total_loss += loss.item()
return total_loss / len(train_loader)
# 4. 評価関数
def evaluate(model, test_loader, criterion, device):
model.eval()
total_loss = 0
with torch.no_grad():
for data, _ in test_loader:
data = data.view(data.size(0), -1).to(device)
output = model(data)
loss = criterion(output, data)
total_loss += loss.item()
return total_loss / len(test_loader)
# 5. 結果の可視化
def visualize_results(model, test_loader, device, num_images=10):
model.eval()
with torch.no_grad():
data, _ = next(iter(test_loader))
data = data[:num_images].to(device)
data_flat = data.view(data.size(0), -1)
# 再構成
reconstructed = model(data_flat)
# 画像を元の形状に戻す
data = data.cpu().view(-1, 28, 28)
reconstructed = reconstructed.cpu().view(-1, 28, 28)
# プロット
fig, axes = plt.subplots(2, num_images, figsize=(15, 3))
for i in range(num_images):
# 元画像
axes[0, i].imshow(data[i], cmap='gray')
axes[0, i].axis('off')
if i == 0:
axes[0, i].set_title('Original', fontsize=10)
# 再構成画像
axes[1, i].imshow(reconstructed[i], cmap='gray')
axes[1, i].axis('off')
if i == 0:
axes[1, i].set_title('Reconstructed', fontsize=10)
plt.tight_layout()
plt.savefig('autoencoder_results.png')
plt.show()
# 6. メイン実行部分
def main():
# ハイパーパラメータ
input_dim = 784 # 28x28
hidden_dim = 128
latent_dim = 32
epochs = 10
learning_rate = 0.001
batch_size = 128
# デバイス設定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
# データ準備
train_loader, test_loader = prepare_data(batch_size)
# モデル初期化
model = BinaryAutoEncoder(input_dim, hidden_dim, latent_dim).to(device)
# 損失関数(二値交差エントロピー)
criterion = nn.BCELoss()
# 最適化手法
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 訓練ループ
print("Training started...")
for epoch in range(1, epochs + 1):
train_loss = train(model, train_loader, optimizer, criterion, device)
test_loss = evaluate(model, test_loader, criterion, device)
print(f'Epoch [{epoch}/{epochs}], '
f'Train Loss: {train_loss:.4f}, '
f'Test Loss: {test_loss:.4f}')
# 結果の可視化
visualize_results(model, test_loader, device)
# モデルの保存
torch.save(model.state_dict(), 'binary_autoencoder.pth')
print("Model saved!")
if __name__ == '__main__':
main()
主要部分の解説
1. モデル構造
Encoder: 784 → 128 → 32 (次元削減)
Decoder: 32 → 128 → 784 (次元復元)
2. 重要なポイント
- Sigmoidの使用: 出力を0-1の範囲に制限
- BCELoss: 二値データに適した損失関数
- 二値化:
(x > 0.5).float()で画像を白黒に変換
3. 損失関数の選択
# 二値交差エントロピー損失
criterion = nn.BCELoss()
このコードをそのまま実行すれば、MNISTータセットで二値画像のAutoEncoderを訓練できます!