AMD GPU ユーザーがGPUを活かして機械学習などを試すには、ROCm を使用するか DirectML を利用するかの選択をする必要があります。

• ROCm
  • Windows でも動作しますが、現状サポートされていない部分も多くあります。ROCm を試す場合は Linux OS で実行したほうがよさそうです。
    • PyTorch on Windows Preview Edition 25.20.01.14 Release 時点で今後サポートが広がる可能性はあるかも。
• DirectML
  • Windows で機械学習を試すならこちらを選択するのが「黄金律」のようです。
  • pyTorchで利用するにはDirectMLをインストールする必要があります。

# pip install torch-directml
import torch_directml

torch_directml.is_available()
torch_directml.device() # id name
torch_directml.device_name(device_id=0) # GPU name
len(torch_directml.gpu_memory(device_id=0)) # 使用可能メモリ
torch_directml.has_float64_support(device_id=0)

概要

画像生成AIモデルにキャラクター（オープンソースのキャラクターitakoさん）を覚えさせてそのキャラクターを生成できるようにする。

• これを実現するには元のモデルをファインチューニングしてキャラクターを覚えさせる方法と、
• モデル自体は変更せず（固定）してキャラクター情報を後付けするようなアダプター（LoRA）を学習する方法とがある。

モデルのファインチューニングは時間もコストもかかり・・一般のPCでは難しい。LoRAなら計算コストを低く実現できるので一般のPCでも作成が可能なので今回はこちらを進めた。

データの用意

データはオープンデータの東北ずん子・ずんだもんプロジェクトさんのitakoさんを利用します。

それに対応するテキストデータを用意します。最初にキーとなるユニークなキーワードを指定し、そのあとにモデルに覚えてもらいたくない部分を指摘したキーワードを追加します。

• iiitako：これを生成時に指定すると今回学習するキャラクターの特徴が現れます。
• miko dress, japanese clothes ：生成時に服装を色々と変えたいので画像の服装は覚えさせません
• white background, simple background

iiitako, 1girl, miko dress, japanese clothes, white background, simple background

LoRA 学習のコアコード

下記の部分がLoRA学習の本質的な学習部分になります。モデル自体は固定してLoRA部分だけを学習していく必要があります。（メモ：この部分は記載していませんが、力業で押し切った感があるので再考の余地あり・・・）

for epoch in range(num_epochs):
    train_loss = 0.0
    pbar = tqdm(dataset, desc=f"Epoch {epoch+1}")
    for batch in pbar: # 画像と対になるテキストを一つずつ抜き出す
        # データセット（画像と対応するテキスト）
        pixel_values = batch["pixel_values"]\
        .unsqueeze(0).to(device, dtype=torch.float32)
        input_ids = batch["input_ids"].unsqueeze(0).to(device)

        # VAE で画像を低次元に圧縮して（潜在空間で表現する）
        # そこから勾配を計算して特徴を抽出する
        with torch.no_grad():
            latents = vae.encode(pixel_values.to(dtype=torch.float32))\
            .latent_dist.sample()
            # VAEの潜在空間スケール調整用定数
            #（標準正規分布からサンプリングされた出力を適切なスケールに調整）
            latents = latents * 0.18215  

        optimizer.zero_grad() # 画像１枚に対して、最初に勾配をリセット
        image_total_loss = 0.0 # logを表示する場合は必要
        
        # 画像１枚に対してstep_per_one_image回の学習ループを回す
        for step_idx in range(step_per_one_image): 
            encoder_hidden_states = text_encoder(input_ids)[0]
            # ノイズ付与：
            #（latentsがノイズを受ける潜在空間内の画像、noiseが追加するノイズ）
            noise = torch.randn_like(latents)
            timesteps = torch.randint(\
            0, noise_scheduler.config.num_train_timesteps, \
            (1,), device=device, dtype=torch.long)

            noisy_latents = noise_scheduler.add_noise(\
            latents, noise, timesteps)

            # 予測：UNetがノイズを予測する（出力はsample属性として取得）
            noise_pred = unet(\
            noisy_latents, timesteps, encoder_hidden_states).sample

            # Loss計算：
            # MSE Lossを使用し、1000倍にスケール（Loss値が小さすぎる場合の調整）
            loss = F.mse_loss(\
            noise_pred.float(), noise.float(), reduction="mean")
            loss = loss * 1000.0

            # 勾配を蓄積：
            # step_per_one_image回分のLossを平均して
            # バックプロパゲーション（勾配の逆伝搬）
            loss = loss / step_per_one_image
            loss.backward()
            image_total_loss += loss.detach().item()

        # 【超重要】UNetとText Encoder両方のLoRA（重み）が更新される
        # 勾配クリップ: 最大勾配ノルムを1.0に制限して爆発的な勾配を防ぐ
        torch.nn.utils.clip_grad_norm_(\
        list(unet.parameters()) + list(text_encoder.parameters()), 1.0)
        optimizer.step()
        lr_scheduler.step()

これを20エポック、画像一枚当たり10回の学習を回して画像を生成してみました。 元になるモデルを学習時とは違うモデルでの生成も可能でイラストの印象をガラッと変えることも可能です。

今回20エポック回したのですが、1エポックごとにLoRAを書き出しながら進めたので、それぞれのLoRAでseed値を固定して同じプロンプトで生成してみました。（左上がepoch 0、右下がepoch 20）最初は特徴をとらえきれていないですが・・・だんだんと猫耳や髪型などの特徴を学習している様子が見て取れます・・目の色に関しては青と紫で迷っている感じがみられます・・実際に上記の生成結果でも青が出たり紫が出たり、時にはオッドアイで出たりもしました。

LoRAを重ねる

今回のitakoさんLoRAに葛飾北斎の「富嶽三十六景」風に出力するLoRAの２つを適用して出力してみた。が2つのLoRAの適用はバランスをとるのが難しく・・葛飾北斎風に寄せるとitakoさんの特徴が無くなり・・itakoさんのLoRAに寄せると北斎風の特徴が無くなっていく・・・・バランスをとるのが非常に難しいと感じたが・・以下はバランスをとってみた画像です

---

[!NOTE] 参照サイト

DirectML

Windows ROCm

東北ずん子・ずんだもんプロジェクト