WSLでPyTorch使う時にGPU使う方法（DirectML）

WSLでPyTorch使う時にGPU使う方法ないか？と調べていたら、nVIDIA系はちらほら情報見かけるけどAMD系がなかなかない（当方AMD環境）。で、たどり着いたのがMicrosoftが作っているDirectMLというもので、環境構築も楽でした。以下、導入方法と使い方。

インストール

2023年3月時点でPyTorchの1.13にしか対応していらしい。

$ pip install torch==1.13 torchvision torch-directml

TensorFlowの場合は、torch-directmlのかわりにtensorflow-directmlを入れればいいだけのようです。

使い方

torch_directmlをインポートしてデバイス指定するだけ。

import torch
import torch_directml

if torch.cuda.is_available():
    device = torch.device('cuda')
elif torch_directml.is_available():
    device = torch_directml.device()
else:
    device = torch.device('cpu')

あとはcuda使う時と同様、 .to(device) を必要なとこにつけてあげましょう。

普通のGPUと、CPU内蔵のGPUも認識してくれるようです。

for n in range(torch_directml.device_count()):
    print(torch_directml.device_name(n))

AMD Radeon RX 6600M
AMD Radeon(TM) Graphics

ちなみに、CUDAを使って直接GPU使うよりDirectMLの方が性能は落ちるそう。海外の検証動画（youtube: Tensorflow CUDA vs DirectML on 3090,Titan RTX and Radeon 6800）によるとDirectMLでトレーニングした場合3倍くらい時間がかかるようです。

というわけで、

• 素のOS環境は汚したくないWSL上で検証したい。
• GPU余らせておくのはもったいない。

という方はお試しあれ。