PytorchでAMD iGPUを利用するのに DirectMLを使うのが一番手っ取り早かったお話（追記：yolov5sが無理やり動いたお話）

ROCｍが対応していない

勢いでぽちっとRyzen AI シリーズMINI PCを買ったのですが、ROCmがまだ対応していなくて悲しい思いをしました。

が。

WSL2上でDirectMLをバックエンドに使ったPytorchを利用できる模様

Pytorch-DirectMLのインストール

pythonの仮想環境を作ってぽちっとインストール

pip install torch-directml

インストールできました。

Python 3.12.3 (main, Feb  4 2025, 14:48:35) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_directml
>>> dml = torch_directml.device()
>>> torch_directml.is_available()
True

本当にGPUつかってくれてるのん？

という疑念もありますのでCPUとDirectMLでベンチマークをいたしますー

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
CPU と DirectML (AMD GPU) で大きな行列の乗算を行い、パフォーマンスを比較するテストコード
"""

import time
import torch
import torch_directml

def benchmark(device, size, iterations=10):
    """
    指定したデバイス上で、size x size の行列乗算を iterations 回実行し、
    1回あたりの平均実行時間（秒）を返す関数
    """
    # ランダムな行列 A, B を作成
    A = torch.randn(size, size, device=device)
    B = torch.randn(size, size, device=device)

    # ウォームアップ: 初回のオーバーヘッドを取り除くために数回実行
    for _ in range(3):
        _ = torch.matmul(A, B)

    times = []
    for i in range(iterations):
        start = time.time()
        C = torch.matmul(A, B)
        # 結果取得することで同期（.item() はブロッキング呼び出し）
        _ = C[0, 0].item()
        elapsed = time.time() - start
        times.append(elapsed)
        print(f"Iteration {i+1}: {elapsed:.6f} seconds")

    average_time = sum(times) / iterations
    return average_time

def main():
    # 行列サイズは大きいほどGPUの並列性が発揮されやすい
    size = 10000  # 行列の次元（必要に応じて調整してください）
    iterations = 10  # 計測する反復回数

    print("【マトリクス乗算ベンチマーク】")
    print(f"行列サイズ: {size}x{size}, 繰り返し回数: {iterations}")

    # --- CPU でのベンチマーク ---
    cpu_device = torch.device("cpu")
    print("\n--- CPU ベンチマーク ---")
    cpu_time = benchmark(cpu_device, size, iterations)
    print(f"CPU 平均実行時間: {cpu_time:.6f} 秒")

    # --- DirectML (AMD GPU) でのベンチマーク ---
    try:
        dml_device = torch_directml.device()
        print("\n--- DirectML (GPU) ベンチマーク ---")
        print("DirectML デバイス:", dml_device)
        dml_time = benchmark(dml_device, size, iterations)
        print(f"DirectML (GPU) 平均実行時間: {dml_time:.6f} 秒")
    except Exception as e:
        print("DirectML デバイスの初期化に失敗しました:", e)

if __name__ == "__main__":
    main()

簡単なテストコードを書きまして……

実行！！

【マトリクス乗算ベンチマーク】
行列サイズ: 10000x10000, 繰り返し回数: 10

--- CPU ベンチマーク ---
Iteration 1: 2.796834 seconds
Iteration 2: 2.767936 seconds
Iteration 3: 2.750078 seconds
Iteration 4: 2.760909 seconds
Iteration 5: 2.690996 seconds
Iteration 6: 2.751591 seconds
Iteration 7: 2.787686 seconds
Iteration 8: 2.576999 seconds
Iteration 9: -0.172003 seconds
Iteration 10: 2.731369 seconds
CPU 平均実行時間: 2.444240 秒

--- DirectML (GPU) ベンチマーク ---
DirectML デバイス: privateuseone:0
Dropped Escape call with ulEscapeCode : 0x03007703
Iteration 1: 6.560529 seconds
Iteration 2: 1.492014 seconds
Iteration 3: 1.502179 seconds
Iteration 4: 1.495882 seconds
Iteration 5: 1.398580 seconds
Iteration 6: 1.506729 seconds
Iteration 7: 1.507093 seconds
Iteration 8: 1.422592 seconds
Iteration 9: 1.506874 seconds
Iteration 10: 1.489866 seconds
DirectML (GPU) 平均実行時間: 1.988234 秒

使ってくれてそう。
この際にWindows側のGPU利用率が上がっているので大丈夫でしょう。
さて、yorov7のpythonコードをちょっと書き直して再学習始めますか……

---

追記

yolov5を落としてきて、train.pyでCUDAを使っている所をスキップしまくってDirectMLデバイスを読み込み、またval.pyでもDirectMLデバイスを読み込むように書き換えた結果。

何とか学習が進みました。
けっこーパッチ当てたのですが、綺麗に書けたら公式にコミットしてもいいのかなあ？
でも一部無理やりが過ぎるのでだめですねー

train.pyを結構書き換え

import argparse
import math
import os
import random
import subprocess
import sys
import time
from copy import deepcopy
from datetime import datetime, timedelta
from pathlib import Path

try:
    import comet_ml  # must be imported before torch (if installed)
except ImportError:
    comet_ml = None

try:
    from torch.cuda.amp import GradScaler as CudaGradScaler
except ImportError:
    CudaGradScaler = None  # CUDAがない環境でもスクリプトが壊れないように

try:
    from torch.amp import GradScaler as CpuGradScaler
except ImportError:
    CpuGradScaler = None  # 古いPyTorchの場合を考慮

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm
import psutil

(snip...)

from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (
    EarlyStopping,
    ModelEMA,
    de_parallel,
    select_device,
    smart_DDP,
    smart_optimizer,
    smart_resume,
    torch_distributed_zero_first,
)

LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))  # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv("RANK", -1))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
GIT_INFO = check_git_info()


try:
    import torch_directml
    device = torch_directml.device()
    is_directml = True
    print("★ DirectML デバイスを使用します")
except ImportError:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    is_directml = False
    print(f"★ 使用デバイス: {device}")


def log_model_device(model):
    devices = set(p.device for p in model.parameters())
    for i, d in enumerate(devices):
        LOGGER.info(f"Model device {i + 1}: {d}")

def train(hyp, opt, device, callbacks):

train関数の中も無理やり。。

    # Config
    plots = not evolve and not opt.noplots  # create plots

    # DirectMLまたは通常のCUDA対応
    if hasattr(torch, "directml"):
        device = torch_directml.device()  # 明示的にDirectMLデバイスを指定
        is_directml = True
        cuda = False  # CUDAではないのでFalseで固定
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        is_directml = False
        cuda = device.type != "cpu"

    init_seeds(opt.seed + 1 + RANK, deterministic=True)

    with torch_distributed_zero_first(LOCAL_RANK):
        data_dict = data_dict or check_dataset(data)  # check if None

    train_path, val_path = data_dict["train"], data_dict["val"]
    nc = 1 if single_cls else int(data_dict["nc"])  # number of classes
    names = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"]  # class names
    is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt")  # COCO dataset

    # Model
    check_suffix(weights, ".pt")  # check weights
    pretrained = weights.endswith(".pt")

    device = torch_directml.device()　# <<< ここは最終手段でむりやり！　本当はdevice判断すべき・・・

    if pretrained:
        with torch_distributed_zero_first(LOCAL_RANK):
            weights = attempt_download(weights)  # download if not found locally

        print(f"★ 使用デバイス: {device}")
        ckpt = torch.load(weights, map_location="cpu")  # load checkpoint to CPU to avoid CUDA memory leak
        print(f"★ 使用デバイス: {device}")
        
        model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create

        log_model_device(model)
        exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else []  # exclude keys
        csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
        csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
        model.load_state_dict(csd, strict=False)  # load
        LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}")  # report
    else:
        model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device)  # create

    amp = check_amp(model)  # check AMP

こんな感じに書き換え。最終手段が無理やりすぎますｗ

あと、val.pyを呼び出しているところはしかたなくCPUに・・・
（2か所あります

            if not noval or final_epoch:  # Calculate mAP
                results, maps, _ = validate.run(
                    data_dict,
                    batch_size=batch_size // WORLD_SIZE * 2,
                    imgsz=imgsz,
                    half=False,
                    model=ema.ema.to("cpu"),  # ここで明示的にCPUに移動
                    device=torch.device("cpu"),  # valはCPUで行う
                    single_cls=single_cls,
                    dataloader=val_loader,
                    save_dir=save_dir,
                    plots=False,
                    callbacks=callbacks,
                    compute_loss=None,    # validateの時は使わない(DirectML対策)
                )

            # Deviceを戻す
            model.to(device)
            ema.ema.to(device)


(snip....)

    # end training -----------------------------------------------------------------------------------------------------
    if RANK in {-1, 0}:
        LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.")
        for f in last, best:
            if f.exists():
                strip_optimizer(f)  # strip optimizers
                if f is best:
                    LOGGER.info(f"\nValidating {f}...")
                    model_cpu = attempt_load(f, map_location="cpu").float()  # 明示的にfloatでロード
                    results, _, _ = validate.run(
                        data_dict,
                        batch_size=batch_size // WORLD_SIZE * 2,
                        imgsz=imgsz,
                        iou_thres=0.65 if is_coco else 0.60,  # best pycocotools at iou 0.65
                        single_cls=single_cls,
                        model=model_cpu,  # ここで明示的にCPUに移動
                        device=torch.device("cpu"),  # valはCPUで行う
                        dataloader=val_loader,
                        save_dir=save_dir,
                        save_json=is_coco,
                        verbose=True,
                        plots=plots,
                        callbacks=callbacks,
                        compute_loss=None,    # validateの時は使わない(DirectML対策)
                    )  # val best model with plots

そうこうしてなんとかうごきましたとさ。
最後まで動くかな・・・