Pythonで学ぶ実際の画像を使ったDiffusion Model入門

import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import matplotlib.pyplot as plt

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

# シンプルなCNNモデル
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 28 * 28, 128)
        self.fc2 = nn.Linear(128, 28 * 28)
    
    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        x = x.view(-1, 1, 28, 28)
        return x

# データセットの準備（MNISTを使用）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

mnist_train = datasets.MNIST(root='./data', train=True, transform=transform, download=True)

# トレーニング用データローダー
train_loader = torch.utils.data.DataLoader(dataset=mnist_train, batch_size=64, shuffle=True)

# モデルの初期化
model = SimpleCNN().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# ノイズの追加
def add_noise(img, noise_factor=0.5):
    noisy_img = img + noise_factor * torch.randn(*img.shape).to(device)
    noisy_img = torch.clamp(noisy_img, 0., 1.)
    return noisy_img

# トレーニング
num_epochs = 5
for epoch in range(num_epochs):
    for images, _ in train_loader:
        images = images.to(device)
        noisy_images = add_noise(images)
        optimizer.zero_grad()
        outputs = model(noisy_images)
        loss = criterion(outputs, images)
        loss.backward()
        optimizer.step()
    
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 結果の表示
def show_images(orig, noisy, denoised):
    orig = orig.cpu().numpy()
    noisy = noisy.cpu().numpy()
    denoised = denoised.cpu().detach().numpy()
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    axes[0].imshow(orig[0][0], cmap='gray')
    axes[0].set_title("Original Image")
    axes[1].imshow(noisy[0][0], cmap='gray')
    axes[1].set_title("Noisy Image")
    axes[2].imshow(denoised[0][0], cmap='gray')
    axes[2].set_title("Denoised Image")
    plt.show()

# サンプル画像の取得と表示
sample_image, _ = mnist_train[0]
sample_image = sample_image.unsqueeze(0).to(device)
noisy_sample = add_noise(sample_image)
denoised_sample = model(noisy_sample)

show_images(sample_image, noisy_sample, denoised_sample)