AnoGANの概要とアルゴリズム及び実装例

pip install torch torchvision matplotlib numpy tqdm

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt

# デバイス設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ハイパーパラメータ
latent_dim = 100
image_size = 28*28
batch_size = 64
epochs = 50
lr = 0.0002

# データセット (正常データ: MNIST の「0」のみ使用)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
mnist = torchvision.datasets.MNIST(root='./data', train=True, transform=transform, download=True)

normal_data = [data for data in mnist if data[1] == 0]  # 数字「0」のみを抽出
train_loader = torch.utils.data.DataLoader(normal_data, batch_size=batch_size, shuffle=True)

# **(1) GAN のモデル定義**
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, image_size),
            nn.Tanh()
        )

    def forward(self, z):
        return self.model(z).view(-1, 1, 28, 28)

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(image_size, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.model(x.view(-1, image_size))

# **(2) モデルの初期化**
G = Generator().to(device)
D = Discriminator().to(device)

# **(3) 損失関数と最適化**
criterion = nn.BCELoss()
optimizer_G = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer_D = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))

# **(4) GAN の学習**
for epoch in range(epochs):
    for i, (images, _) in enumerate(train_loader):
        images = images.view(-1, image_size).to(device)
        batch_size = images.size(0)

        # ラベル
        real_labels = torch.ones(batch_size, 1).to(device)
        fake_labels = torch.zeros(batch_size, 1).to(device)

        # **Discriminator の学習**
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_images = G(z).detach()
        loss_D = criterion(D(images), real_labels) + criterion(D(fake_images), fake_labels)

        optimizer_D.zero_grad()
        loss_D.backward()
        optimizer_D.step()

        # **Generator の学習**
        z = torch.randn(batch_size, latent_dim).to(device)
        fake_images = G(z)
        loss_G = criterion(D(fake_images), real_labels)  # Generator の目的は本物らしい画像を生成すること

        optimizer_G.zero_grad()
        loss_G.backward()
        optimizer_G.step()

    print(f"Epoch [{epoch+1}/{epochs}] Loss_D: {loss_D.item():.4f} Loss_G: {loss_G.item():.4f}")

print("Training Finished!")

# **(5) 異常スコアの計算**
def anomaly_score(x, G, z_dim=100, steps=500, lr=0.1):
    z = torch.randn(1, z_dim, requires_grad=True, device=device)
    optimizer = torch.optim.Adam([z], lr=lr)

    for _ in range(steps):
        fake_x = G(z)
        loss = torch.mean((fake_x - x) ** 2)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    return loss.item()

# **(6) 異常データの検出**
test_data = torchvision.datasets.MNIST(root='./data', train=False, transform=transform, download=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)

G.eval()
anomaly_scores = []

for i, (img, label) in enumerate(test_loader):
    img = img.to(device)
    score = anomaly_score(img, G)
    anomaly_scores.append((score, label.item()))

    if i >= 100:
        break

# **(7) 結果の可視化**
scores, labels = zip(*anomaly_scores)
plt.hist(scores, bins=30, alpha=0.7, label="Anomaly Score")
plt.xlabel("Anomaly Score")
plt.ylabel("Frequency")
plt.legend()
plt.show()