PyTorchによる二値画像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()

Encoder: 784 → 128 → 32 (次元削減)
Decoder: 32 → 128 → 784 (次元復元)

# 二値交差エントロピー損失
criterion = nn.BCELoss()

$ python3 binary-image-autoender.py
Using device: cpu
Training started...
Epoch [1/10], Train Loss: 0.2104, Test Loss: 0.1398
Epoch [2/10], Train Loss: 0.1216, Test Loss: 0.1047
Epoch [3/10], Train Loss: 0.1003, Test Loss: 0.0925
Epoch [4/10], Train Loss: 0.0900, Test Loss: 0.0841
Epoch [5/10], Train Loss: 0.0830, Test Loss: 0.0786
Epoch [6/10], Train Loss: 0.0785, Test Loss: 0.0753
Epoch [7/10], Train Loss: 0.0754, Test Loss: 0.0725
Epoch [8/10], Train Loss: 0.0731, Test Loss: 0.0707
Epoch [9/10], Train Loss: 0.0714, Test Loss: 0.0694
Epoch [10/10], Train Loss: 0.0701, Test Loss: 0.0682
Model saved!