MENU
  1. ホーム
  2. ナレッジ
  3. 【Google Colab】64×64の花画像で作る拡散モデル

【Google Colab】64×64の花画像で作る拡散モデル

ナレッジ
AI Python 拡散モデル 画像生成
目次

1. はじめに：「花を復元するAI」

まずはこれを見てください。

完全なランダムノイズ（砂嵐）から、1000ステップかけて花が浮かび上がってくる様子

最初のフレームは、ただの砂嵐です。意味のある情報は1ビットも入っていません。それが数百ステップ進むと、なんとなく色の塊が現れ、輪郭がまとまり、最後には花のような物体が描かれます。

この画像を生成しているのは、あなたがこれから書き上げるコードです。事前学習済みモデルのダウンロードではありません。ノイズスケジュールの定義から学習ループ、サンプリングまで、全部自分の手で書きます。

動作環境

  • Google Colab（無料枠のT4 GPUでOK）
  • ランタイム → ランタイムのタイプを変更 → ハードウェアアクセラレータ → T4 GPU を忘れずに設定する。
  • 今回は学習を複数回に分けて再開できる構成にしています。Colabのセッションが12時間で切れても、続きから学習を再開できます（詳しくは6章）

2. 拡散モデルの全体像

数式の前に、図で理解しましょう。拡散モデルがやっていることは、驚くほど単純な2つのプロセスだけです。

順伝播（Forward Process）：花を砂嵐にする

きれいな花の画像に、ほんの少しだけノイズを足します。ほとんど変化はありません。もう一度足します。まだ花です。これを1000回繰り返すと、元が何だったのか完全にわからない砂嵐になります。

ここで重要なのは、このプロセスに学習は一切必要ないということです。「乱数を足す」だけなので、機械的に実行できます。プログラムで言えば image + random_noise です。

逆伝播（Reverse Process）：砂嵐から花を取り戻す

次は逆伝播（ノイズを取り除く作業）です。砂嵐から少しずつノイズを引き算していけば、花に戻るはずです。

人間にとっては「何を引けばいいのか」がわかりません。

ここでAIの出番です。

AIが学習する唯一のこと

拡散モデルのニューラルネットワークが学習するのは、次の１タスクのみです。

「ノイズまみれの画像から、そこに乗っている『ノイズ成分』だけを当てる」

花の描き方を学ぶわけではなく、ノイズを推定するだけです。そして推定したノイズを引き算すれば、結果として花が残るということになっています。

しかも、この学習には正解ラベルを準備する必要がありません。順伝播で自分が足したノイズが正解データになるからです。

一言でまとめると、拡散モデルの学習とは：

「自分で仕込んだノイズを、自分で当てる訓練を延々と繰り返す」

となっています。以降の章では、以上の流れを数式とコードで記述していきます。

3. 実装準備：Colab環境とデータセット

ライブラリのインストール

!pip install -q diffusers datasets
import torch<br>import torch.nn.functional as F
from torch.optim import Adam
from diffusers import UNet2DModel
from datasets import load_dataset
from torchvision import transforms
import os

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)  # cuda と出ればOK

データセットのロード

Hugging Face Hubから花の画像データセットを取ってきます。

dataset = load_dataset("huggan/flowers-102-categories", split="train")
print(dataset)
print(dataset.column_names)  # 画像カラムの名前を確認

huggan/flowers-102-categories は、102種類の花を集めた画像データセットです。

前処理：64×\times64へのリサイズとテンソル化

preprocess = transforms.Compose([
    transforms.Resize((64, 64)),
    transforms.ToTensor(),                  # [0, 1] のテンソルへ
    transforms.Normalize([0.5], [0.5]),     # [-1, 1] へ
])

def transform(examples):
    return {"image": [preprocess(img.convert("RGB")) for img in examples["image"]]}

dataset.set_transform(transform)

dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)

ここで地味だが極めて重要なのが Normalize([0.5], [0.5]) です。画素値を [0, 1] ではなく [-1, 1] の範囲に変換しています。

理由は、この後で足す標準正規分布のノイズ 𝒩(0,1)\mathcal{N}(0, 1) が、平均0で正負に広がる値だからです。画像側が [0, 1] のままだと、データとノイズのスケールが噛み合わず、学習がうまく進みません。

（[0.5] と1要素しか渡していませんが、3チャンネルの画像に対してもブロードキャストされるので、実質 [0.5, 0.5, 0.5] と同じ意味になります。）

set_transform を使うと、examples["image"] のキー名で辞書を返す必要があります。あとで学習ループ側は batch["image"] としてこれを受け取ります。

中身の確認

import matplotlib.pyplot as plt

batch = next(iter(dataloader))["image"]
print(batch.shape)  # torch.Size([16, 3, 64, 64])

fig, axes = plt.subplots(1, 8, figsize=(16, 2))
for ax, img in zip(axes, batch[:8]):
		# Matplotlibのimshowは(高さ, 幅, チャンネル数)の順序を期待するため、
    # テンソルを(チャンネル数, 高さ, 幅)から(高さ, 幅, チャンネル数)に並べ替える必要がある。
    # permute(1, 2, 0)は、元のインデックス1(高さ)を新しいインデックス0に、
    # 元のインデックス2(幅)を新しいインデックス1に、
    # 元のインデックス0(チャンネル数)を新しいインデックス2に移動させる。
    ax.imshow((img.permute(1, 2, 0) + 1) / 2)  # [-1,1] → [0,1] に戻して表示
    ax.axis("off")
plt.show()

花が8輪並べば準備完了です。

4. 順伝播：画像を砂嵐にする数式とコード

数式の工夫：1000回のループを1行にする

素朴に実装するなら、「ノイズを足す」を tt 回ループすれば、時刻 ttのノイズ画像 xtx_t が手に入ります。しかし学習中に毎回1000回ループを回すのは論外です。

そこでDDPMの論文が使うのが、任意の時刻 tt のノイズ画像を一発で作れるこの式です。

q(xt|x0)=𝒩(xt;α‾tx0,(1−α‾t)𝐈)q(x_t \vert x_0) = \mathcal{N}\left(x_t; \sqrt{\bar{\alpha}_t}x_0, (1-\bar{\alpha}_t)\mathbf{I}\right)

𝒩\mathcal{N} という記号にひるまないでください。これは「正規分布」を表すだけで、意味しているのは次のことです。

xtx_t は、平均が α‾tx0\sqrt{\bar{\alpha}_t}x_0、分散が 1−α‾t1-\bar{\alpha}_t の正規分布から取り出される

そして正規分布からのサンプリングは、「平均 + 標準偏差 × 標準正規乱数」で書けます。つまり上の式は、実質こう読み替えられます。

xt=α‾tx0+1−α‾tϵ(ϵ∼𝒩(0,𝐈))x_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon \qquad (\epsilon \sim \mathcal{N}(0, \mathbf{I}))

元画像 x0x_0 とノイズ ϵ\epsilon の、重み付き足し算を表しています。

  • α‾t\sqrt{\bar{\alpha}_t}：元画像を残す割合（tt が進むほど0に近づく）
  • 1−α‾t\sqrt{1-\bar{\alpha}_t} ：ノイズを混ぜる割合（tt が進むほど1に近づく）

t=0t=0 では画像100%、t=1000t=1000 ではノイズ100%になるはずです。

β\beta、α\alpha、α‾\bar{\alpha} の関係

登場する記号は3つだけで、後ろ2つは β\beta から機械的に決まります。

記号定義意味
β\beta自分で決める（スケジュール）各ステップで足すノイズの強さ
αt\alpha_t1−βt1 – \beta_t各ステップで残る画像の割合
α‾\bar{\alpha}α1α2⋯αt\alpha_1 \alpha_2 \cdots \alpha_tステップ0から tt まで残った画像の割合（累積積）

α‾\bar{\alpha} は「α\alpha の累積積」というだけで、PyTorchでは torch.cumprod により実装できます。

PyTorch実装

T = 1000  # 総ステップ数

# 線形ノイズスケジュール（DDPM論文の設定）
beta = torch.linspace(0.0001, 0.02, T).to(device)
alpha = 1.0 - beta
alpha_cumprod = torch.cumprod(alpha, axis=0).to(device)  # ← これが ᾱ_t

そして、順伝播の本体です。

def q_sample(x_0, t, noise=None):
    """x_0 に時刻 t 相当のノイズを足す"""
    if noise is None:
        noise = torch.randn_like(x_0)

    sqrt_alpha_cumprod_t = torch.sqrt(alpha_cumprod[t])[:, None, None, None]
    sqrt_one_minus_alpha_cumprod_t = torch.sqrt(1.0 - alpha_cumprod[t])[:, None, None, None]

    return sqrt_alpha_cumprod_t * x_0 + sqrt_one_minus_alpha_cumprod_t * noise

alpha_cumprod[t] は形が (batch_size,) のベクトルですが、画像は (batch_size, 3, 64, 64) という4次元テンソルです。[:, None, None, None] で末尾に3つの次元を追加し、(batch_size, 1, 1, 1) の形にしてからブロードキャストしています。

最後の1行に注目してください。 上で紹介した数式

xt=α‾tx0+1−α‾tϵx_t = \sqrt{\bar{\alpha}_t}x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon

が、そのまま sqrt_alpha_cumprod_t * x_0 + sqrt_one_minus_alpha_cumprod_t * noise になっています。記号がPyTorchの変数名に置き換わっただけです。

実際に花にノイズを加えてみる

x_0 = batch[:1].to(device)  # 花を1輪
steps = [0, 50, 100, 200, 400, 700, 999]

fig, axes = plt.subplots(1, len(steps), figsize=(16, 2.5))
for ax, step in zip(axes, steps):
    t = torch.tensor([step], device=device)
    x_t = q_sample(x_0, t)
    img = (x_t[0].cpu().permute(1, 2, 0) + 1) / 2
    ax.imshow(img.clamp(0, 1))
    ax.set_title(f"t={step}")
    ax.axis("off")
plt.show()

出力を見ると、t=50 あたりではまだ花らしさが残っていて、t=200 でざらつきが目立ち、t=400 で輪郭がぼやけ、t=999 では完全な砂嵐になっています。

5. モデル定義：ノイズを予測する「U-Net」

U-Netの役割は「画像 → 画像」

拡散モデルのネットワークがやることは、非常にシンプルです。

  • 入力：ノイズまみれの画像 xtx_t（3×64×643 \times 64 \times 64）と、時刻 tt（整数1つ）
  • 出力：そこに乗っているノイズ ϵ\epsilon の予測値（3×64×643 \times 64 \times 64）

入力と出力が同じサイズです。分類のように次元を潰していく必要はなく、画像を入れたら画像が出てくる構造が必要になります。

そこでU-Netです。名前の由来である「U」の字は、こういう形をしています。

前半（Encoder）で解像度を落としながら「全体の構造」を捉え、後半（Decoder）で解像度を戻します。このとき、同じ解像度の層どうしを横に直結する「スキップ接続」があるのがU-Netの肝です。ノイズ予測は「1ピクセル単位の細かい値」を推定するタスクです。8×88\times 8まで潰した情報だけから64×6464\times 64を復元しようとすると、細部が完全に失われます。スキップ接続は、Encoderが持っていた高解像度の情報を、Decoderに直接手渡す抜け道になっています。

Time Embedding

同じネットワークが、t=10t=10 のときと t=900t=900 のときで、まったく違う振る舞いをしなければなりません。

t=10t=10 の画像はほぼ元画像で、乗っているノイズはごくわずか。t=900t=900 の画像はほぼ砂嵐で、ほとんどがノイズ。同じ「ノイズを当てろ」という仕事でも、要求される出力のスケールも性質もまるで違います。

だからモデルには、「今が何ステップ目なのか」を必ず教えなければなりません。整数 tt をそのまま渡しても、ニューラルネットワークは大きなスカラー1個から細かい違いを読み取るのが苦手です。900と901の差も、10と900の差も、扱いが雑になってしまいます。

そこで使われるのが、TransformerでおなじみのSinusoidal Position Embeddings（正弦波位置エンコーディング）です。整数 tt を、周波数の異なるsin/cosの組み合わせで多次元のベクトルに展開します。

コードの工夫：U-Netはdiffusersから

U-Netの中身は既製品を使います。Hugging Faceの diffusers にある UNet2DModel が、まさにこの用途のために用意されています。

model = UNet2DModel(
    sample_size=64,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(64, 128, 128, 256),
    down_block_types=(
        "DownBlock2D",      # 64x64 → 32x32
        "DownBlock2D",      # 32x32 → 16x16
        "AttnDownBlock2D",  # 16x16 → 8x8（ここでAttention）
        "DownBlock2D",
    ),
    up_block_types=(
        "UpBlock2D",
        "AttnUpBlock2D",    # 16x16でAttention
        "UpBlock2D",
        "UpBlock2D",
    ),
)

model.to(device)

n_params = sum(p.numel() for p in model.parameters())
print(f"パラメータ数: {n_params / 1e6:.1f}M")

Attentionを16×1616\times 16の解像度にだけ入れているのは、計算量とのトレードオフです。Attentionは解像度の2乗でコストが増えるため、64×6464\times 64 に入れると学習に時間がかかりすぎます。一方、低解像度側に入れると「花びらどうしの位置関係」のような大域的な整合性を取るのに効果的です。

呼び出し方は以下の通りです。

t = torch.randint(0, T, (16,), device=device).long()
x_t = q_sample(batch.to(device), t)

pred_noise = model(x_t, t).sample   # ← .sample を忘れずに
print(pred_noise.shape)             # torch.Size([16, 3, 64, 64])

diffusers のモデルは出力をオブジェクトで返すので、テンソルを取り出すには .sample が必要です。ここは地味にハマるポイントです。

なお、時刻の埋め込みは UNet2DModel の内部で自動的に処理されます。前章で説明したSinusoidal Position Embeddingsと本質的に同じものが中で動いている、と思ってください。

6. 学習ループ：ノイズ推定

損失関数

拡散モデルの論文には、変分下界だのKLダイバージェンスだのが数ページにわたって展開されています。しかし、それを全部整理して簡略化した結果、実装で使う損失関数はこうなります。

L=𝔼t,x0,ϵ[‖ϵ−ϵθ(xt,t)‖2]L = \mathbb{E}_{t, x_0, \epsilon} \left[ \left\Vert \epsilon – \epsilon_\theta(x_t, t) \right\Vert^2 \right]

𝔼\mathbb{E} は「平均」、‖⋅‖2\Vert \cdot \Vert^2 は「二乗」。つまり中身は

本物のノイズ ϵ\epsilon と、モデルが予測したノイズ ϵθ\epsilon_\theta の、二乗誤差の平均

したがってこれはただの平均二乗誤差（MSE）です。

loss = F.mse_loss(noise_pred, noise)

コードの工夫：チェックポイントから学習を再開できるようにする

Colabの無料GPUには、セッションが一定時間で切れるという制約があります。花のような被写体は、それなりのエポック数を積まないと品質が出ません。そこで今回は、「途中まで学習したモデルを保存し、次のセッションでそこから続きを学習する」という運用を組み込みます

  1. 学習開始時に、もし保存済みの重み（model.pth）があれば読み込む
  2. 追加で学習を回す
  3. 一定エポックごとに新しい重み（model.pth）として保存し直す

「同じモデルの学習を、セッションをまたいで継続する」といった操作です。

model = UNet2DModel(
    sample_size=64,
    in_channels=3,
    out_channels=3,
    layers_per_block=2,
    block_out_channels=(64, 128, 128, 256),
    down_block_types=("DownBlock2D", "DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
    up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D", "UpBlock2D"),
)

# 前回のセッションで保存した重みがあれば読み込む
if os.path.exists('model.pth'):
    model.load_state_dict(torch.load('model.pth', map_location=device))
    print("Successfully reloaded model.pth (Original Structure)")

optimizer = Adam(model.parameters(), lr=5e-5)
epochs = 50

ここで注意したいのが学習率です。 ゼロから学習するときの 1e-4 に対して、再開後の学習率は 5e-5 と半分に落としています。これは、すでにある程度学習が進んだモデルに対して大きすぎる学習率をかけると、せっかく学んだ重みを一気に壊してしまうことがあるためです。

学習ループ本体

print(f"Fine-tuning started (Saving to model.pth)...")
try:
    for epoch in range(epochs):
        loss_sum = 0
        for step, batch in enumerate(dataloader):
            clean_images = batch["image"].to(device)
            t = torch.randint(0, T, (clean_images.shape[0],), device=device).long()
            noise = torch.randn_like(clean_images).to(device)
            noisy_images = q_sample(clean_images, t, noise=noise)

            noise_pred = model(noisy_images, t).sample
            loss = F.mse_loss(noise_pred, noise)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            loss_sum += loss.item()

        if (epoch + 1) % 5 == 0:
            avg_loss = loss_sum / len(dataloader)
            print(f"Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f}")
            torch.save(model.state_dict(), 'model.pth')

    torch.save(model.state_dict(), 'model.pth')
    print("Final model saved as model.pth")
except Exception as e:
    print(f"Error: {e}")

ループの中核は、実質4ステップだけです。

  1. ランダム値の割り振り：各画像にランダムな tt を割り当てる
  2. ノイズ合成：ノイズを生成して画像に混ぜる
  3. 答え合わせ：モデルの予測と本物のノイズを比較
  4. 5エポックごとに保存：model.pth として重みを更新していく

try/except で囲っているのは、Colabのセッション切断やGPUメモリ不足で学習が途中で止まってしまっても、それまでに保存された model.pth は残るようにするためです。

損失の値の読み方

「lossが下がりきったように見えても、すぐに収束したと判断するのは早計です。 拡散モデルのlossは全時刻の平均値であり、その大部分を「簡単な問題」が占めています。tt が大きい領域（ほぼ砂嵐）では、「入力とほぼ同じものを生成する」のでAIはそこそこに正解できてしまうからです。

一方、生成品質を決めるのは tt が小さい領域の、細かい精度です。この部分の改善については、lossの非常に小さいオーダーにしか現れません。したがってlossの数値ではなく、実際に生成した画像を見て判断する必要があると言えます。

7. 生成： 砂嵐から花を錬成する（サンプリング）

学習が終わるといよいよ砂嵐から花を取り出す段階です。ここからは学習スクリプトとは別のセルとして、生成専用のコードを組みます。

逆プロセスの数式

xt−1=1αt(xt−1−αt1−α‾tϵθ(xt,t))+σtzx_{t-1} = \frac{1}{\sqrt{\alpha_t}} \left( x_t – \frac{1 – \alpha_t}{\sqrt{1 – \bar{\alpha}t}} \epsilon_\theta(x_t, t) \right) + \sigma_t z

やっていることは3手順です。

  1. 括弧の中：現在の画像 xtx_t から、モデルが予測したノイズ ϵθ\epsilon_\theta を、係数をかけて引く
  2. 1/αt1/\sqrt{\alpha_t}：引き算で縮んだ分のスケールを戻す
  3. +σtz+ \sigma_t z：新しいノイズを少しだけ足し戻す

1と2は納得できるとして、3が奇妙に見えるはずです。ノイズを消したいのに、なぜ足すのでしょうか。

理由は、逆プロセスが本質的に確率的だからです。ある砂嵐から復元しうる花は1輪ではなく、無数にあります。毎ステップをノイズを含まない最もらしいものに決め打ちしてしまうと、生成される画像はいずれも均一的な画像に収束します。少しの乱数を残すことで、モデルは毎回違う花を、はっきりした特徴とともに描くことができます。ただし最終ステップ（t=0t=0）では、出力を確定させるためにノイズを足しません。

再現性の確保：シード固定

生成結果を人に見せたり、あとで見比べたりするときは、乱数のシードを固定しておくと便利です。同じシードなら同じ砂嵐から始まるので、モデルの重みを差し替えたときの品質の違いをフェアに比較できます。

import random
import numpy as np

def set_seed(seed=42):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

set_seed(43)

実装

from tqdm.auto import tqdm

model_path = 'model.pth'
if not os.path.exists(model_path):
    model_path = 'model.pth'

model.load_state_dict(torch.load(model_path, map_location=device))
model.to(device)
model.eval()

img = torch.randn(1, 3, 64, 64).to(device)
step_images = []

with torch.no_grad():
    for i in tqdm(reversed(range(T)), total=T):
        t = torch.tensor([i], device=device).long()
        predicted_noise = model(img, t).sample

        a_t, a_cp_t = alpha[i], alpha_cumprod[i]
        img = (1 / torch.sqrt(a_t)) * (
            img - ((1 - a_t) / torch.sqrt(1 - a_cp_t)) * predicted_noise
        )

        if i > 0:
            b_t = beta[i]
            img = img + torch.sqrt(b_t) * torch.randn_like(img)

        img = torch.clamp(img, -1.0, 1.0)

        if i % 200 == 0 or i == 0:
            step_images.append(img.detach().cpu())

数式との対応を確認してください。

数式コード
ϵθ(xt,t)\epsilon_\theta(x_t, t)predicted_noise
xt−1−αt1−α‾tϵθx_t – \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}t}} \epsilon_\thetaimg - ((1 - a_t) / torch.sqrt(1 - a_cp_t)) * predicted_noise
1αt(⋯)\frac{1}{\sqrt{\alpha_t}}(\cdots)(1 / torch.sqrt(a_t)) * (...)
+σtz+\sigma_t z+ torch.sqrt(b_t) * torch.randn_like(img)（σt=βt\sigma_t = \sqrt{\beta_t}）

torch.clamp(img, -1.0, 1.0) について補足すると、理論上、画像は [-1, 1] の範囲に収まっているはずですが、実際にはステップを重ねる中でモデルの予測誤差が蓄積し、値がその範囲をわずかにはみ出すことがあります。毎ステップの終わりにクランプ（値の切り詰め）を入れておくことで、誤差が次のステップに増幅して伝わるのを防ぎ、生成を安定させています。

実行と保存

final_img_tensor = step_images[-1].squeeze().permute(1, 2, 0)
final_img_np = ((final_img_tensor.numpy() + 1.0) / 2.0 * 255).astype(np.uint8)

from PIL import Image
pil_img = Image.fromarray(final_img_np)
pil_img.save("generated_output.png")
print("✅ 画像を generated_output.png として保存しました。")

fig, axes = plt.subplots(1, len(step_images), figsize=(15, 3))
for ax, img_tensor in zip(axes, step_images):
    display_img = (img_tensor.squeeze().permute(1, 2, 0).numpy() + 1.0) / 2.0
    ax.imshow(np.clip(display_img, 0, 1))
    ax.axis("off")
plt.show()

i % 200 == 0 の間隔で保存しているので、step_images には t=800,600,400,200,0t=800, 600, 400, 200, 0 相当の5枚が並びます。砂嵐から花が姿を現すまでを、5コマの遷移図として確認できるはずです。

8. まとめと注意点

何ができたか

64×6464\times 64という低解像度、小さなモデル、無料のT4 GPUで学習。この条件でも、乱数から花が生成されるところまで確認できました。

生成結果がイマイチだったら

改善の効きやすい順に挙げます。

  1. 追加学習を繰り返す — model.pth を次回以降の model.pth として使い回し、さらにエポックを重ねる。
  2. EMA（指数移動平均）を導入する — モデルの重みの移動平均でサンプリングすると、生成が安定する。
  3. モデルを大きくする — block_out_channels を (128, 128, 256, 256) などに
  4. バッチサイズを上げる — 花のように被写体のばらつきが大きいデータでは、batch_size=16 はやや小さめです。GPUメモリに余裕があれば32〜64に増やす。
CTA
ナレッジ
AI Python 拡散モデル 画像生成
  • URLをコピーしました！
  • URLをコピーしました！
この記事を書いた人
上野優太のアバター 上野優太
関連記事
目次