Pythonで学ぶ簡単なDiffusion Model入門

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

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

# シンプルなニューラルネットワークモデル
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(2, 128),
            nn.ReLU(),
            nn.Linear(128, 2)
        )
    
    def forward(self, x):
        return self.fc(x)

# データセットの生成
def generate_data(num_samples=1000):
    angles = np.linspace(0, 2 * np.pi, num_samples)
    x = np.sin(angles)
    y = np.cos(angles)
    return np.stack([x, y], axis=1)

# ノイズの追加
def add_noise(data, noise_level):
    noise = np.random.normal(scale=noise_level, size=data.shape)
    return data + noise

# トレーニングデータの作成
data = generate_data(1000)
noisy_data = add_noise(data, noise_level=0.1)

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

# トレーニング
num_epochs = 1000
for epoch in range(num_epochs):
    optimizer.zero_grad()
    data_torch = torch.tensor(noisy_data, dtype=torch.float32).to(device)
    target_torch = torch.tensor(data, dtype=torch.float32).to(device)
    output = model(data_torch)
    loss = criterion(output, target_torch)
    loss.backward()
    optimizer.step()
    if (epoch + 1) % 100 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 結果のプロット
with torch.no_grad():
    denoised_data = model(data_torch).cpu().numpy()

plt.figure(figsize=(8, 4))
plt.subplot(1, 2, 1)
plt.scatter(noisy_data[:, 0], noisy_data[:, 1], color='red', label='Noisy Data')
plt.scatter(data[:, 0], data[:, 1], color='blue', label='Original Data')
plt.legend()

plt.subplot(1, 2, 2)
plt.scatter(denoised_data[:, 0], denoised_data[:, 1], color='green', label='Denoised Data')
plt.scatter(data[:, 0], data[:, 1], color='blue', label='Original Data')
plt.legend()

plt.show()