AMD Radeon GPUでTensorFlowを動かす - tensorflow-directml編 #TensorFlow

はじめに

TensorFlowの公式では、CUDAベース、つまりNVIDIAのGPU向けでの利用が記載されており、AMD GPU向けとはなっていない。しかし、AMDであろうとせっかくGPUがあるのだから機械学習に使ってみたい！ということで今回はtensorflow-directmlの紹介です。

AMDとしても機械学習利用を広めていきたいということで、ROCmライブラリーというのを出している（TensorFlowとPyTorchが使える）。
https://www.amd.com/ja/graphics/servers-solutions-rocm-ml
（正直、あまりメジャーではない印象ですが、実際どうなのでしょうか。。。今回はこちらの紹介ではありません。）

tensorflow-directml とは

さて、本題の背景です。Microsoftとしても機械学習を推進しており、その中にはDirect Machine Learning (DirectML)というDirectXスタイルで動作するlow-levelのAPIを整備をしているとのこと。

正直DirectX系は触れたことがなく技術の中身についてはわからないのですが、つまるところDirectXに計算を流すことでGPUを動作させるというイメージで、DirectXサポートしたGPUであれば利用できる（詳しい方、詳細な解説など頂ければ。）

上記ページ内のGPU利用のケースを見ていくと、TensorFlow with DirectML on Windowsという項目があるので、本稿ではこれをベースに進める。今回はScoopでパッケージ管理しながらやっていきます。

Scoopのインストール

PowerShellより下記を実行して、Scoopをインストールします。
いきなり２行目を実行し、エラーが出たときに1行目を実行すれば良いです。エラーはポリシーに関するメッセージで、１行目が解消のためのポリシー設定変更になります。実行後確認がでるので「Y」を入力して先に進めましょう。

Set-ExecutionPolicy RemoteSigned -scope CurrentUser
invoke-Expression (New-Object System.Net.WebClient).DownloadString('https://get.scoop.sh')

scoop helpもしくはscoop -vと入力して、ちゃんと出力されればインストールOKです。
バージョン確認scoop -vではGitがインストールされていないとエラーが出ます。今後使うのでインストールしておきます。

scoop install git

実行後はC:\<dir>\scoop\apps\にインストールされます。

Pythonのインストール

Pythonをインストールしていきますが、下記tensorflow-directmlの公式を確認するとPython3.8以降は非サポートと記載があるため、3.7をインストールする必要があります。

scoop bucket add versions #過去のバージョンをインストールするため
scoop install python37 #python3.7の最新版インストール。

C:\<dir>\scoop\apps\python37\3.7.9\ にインストールされます。
python37と実行し、インタプリタが動作すればインストールOKです。

# 以下出力
# Python 3.7.9 (tags/v3.7.9:xxxxx) [] on win32
# xxxxx
exit()  #インタプリタモード終了

ちなみに、自環境ではPython3.10がインストール済でしたが、実行後Path優先度が変更されてしまい、 `python` だとPython3.7.9がデフォルトで動作するようになりました。Path環境は要注意です。

仮想環境の構築

virtualenvを使って、仮想環境に構築していきます。

pip install virtualenv #virtualenvのインストール
mkdir directml #directmlという仮想環境用のフォルダを作成。適宜フォルダパス変更。
virtualenv -p "C:\<dir>\scoop\apps\python37\3.7.9\python.exe" ./directml #directml内に仮想環境を構築する。pはインタプリタを指定するオプション。

仮想環境の開始

.\directml\Scripts\activate #適宜フォルダパス変更。

以降は、powershellの実行環境が(directml)PS C:\<dir>\と仮想実行環境に切り替わっていること。
この仮想環境上ではpythonと実行すると、Python 3.7.9が動作します。

tensorflow-directmlのインストールとデバイス確認

pipで簡単にtensorflow-directmlをインストールできる。

pip install tensorflow-directml  #tensorflow-directmlをインストール

python上で、TensorFlowを使ったGPUデバイス確認を実行してみる。

from tensorflow.python.client import device_lib # GPUデバイス情報取得ライブラリのインポート
device_lib.list_local_devices() # GPUデバイス情報取得

以降出力。時刻他、適当に変換しています。

# 2021-11-10 21:21:21.212121: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2
# 2021-11-10 21:21:21.212121: I tensorflow/stream_executor/platform/default/dso_loader.cc:97] Successfully opened dynamic library C:\<dir>\userlocaldir\directml\lib\site-packages\tensorflow_core\python/directml.24bfac66e4ee42ec393a5fb471412d0177bc7bcf.dll
# 2021-11-10 21:21:21.212121: I tensorflow/stream_executor/platform/default/dso_loader.cc:97] Successfully opened dynamic library dxgi.dll
# 2021-11-10 21:21:21.212121: I tensorflow/stream_executor/platform/default/dso_loader.cc:97] Successfully opened dynamic library d3d12.dll
# 2021-11-10 21:21:21.212121: I tensorflow/core/common_runtime/dml/dml_device_cache.cc:250] DirectML device enumeration: found 1 compatible adapters.
# 2021-11-10 21:21:21.212121: I tensorflow/core/common_runtime/dml/dml_device_cache.cc:186] DirectML: creating device on adapter 0 (AMD Radeon RX 6600 XT)
# 2021-11-10 21:21:21.212121: I tensorflow/stream_executor/platform/default/dso_loader.cc:97] Successfully opened dynamic library Kernel32.dll
# [name: "/device:CPU:0"
# device_type: "CPU"
# memory_limit: 268435456
# locality {
# }
# incarnation: 5232323232323232323
# , name: "/device:DML:0"
# device_type: "DML"
# memory_limit: 7351173325
# locality {
# }
# incarnation: 6802323232323232323
# physical_device_desc: "{\"name\": \"AMD Radeon RX 6600 XT\", \"vendor_id\": 4098, \"device_id\": 23232, \"driver_version\": \"30.0.13025.5005\"}"
# ]

最後のところでAMD Radeon RX 6600 XTとあり、無事GPU情報が取得できていそうです。

ここでの注目はdevice_type: "DML"かなと思います。
実際にTensorFlow使って学習するときには、GPUではなくDMLと表示されます。

tensorflow-directmlでの計算実行

MSのページではこれで終わっていますが、最後に簡単な計算を実行してみます。

import tensorflow.compat.v1 as tf
tf.enable_eager_execution(tf.ConfigProto(log_device_placement=True))

print("TensorFlow ver.=", tf.__version__) #TensorFlowのバージョン確認
print("GPU available? ", tf.test.is_gpu_available()) #GPU利用可否確認

print(tf.add([1.0, 2.0], [3.0, 4.0]))

#以降出力結果。print文とデバイスを示す部分のみ。
TensorFlow ver.= 1.15.5
# 2021-11-10 21:21:21.212121: I tensorflow/core/common_runtime/eager/execute.cc:571] Executing op Add in device /job:localhost/replica:0/task:0/device:DML:0
GPU available?  True
tf.Tensor([4. 6.], shape=(2,), dtype=float32)

tensorflow-directml はv1系TensorFlowのみサポート (2021年11月現在)。バージョン1.15.5がインストールされています。
DMLを利用して、ちゃんとTensorFlowで計算されています。

終了

exit()でインタプリタ終了。
deactivateで仮想環境終了。

以上です。
思ったよりも簡単で助かりました。実際は、DirectX経由するためか遅いなどの問題もあるようですが、githubにはIssueとしても上がっており、改善していくでしょう。
最新版のTensorFlowは利用できかったりと色々とデメリットはありますが、今後に期待していこうと思います。

参考リンク

Microsoft windows AI

Microsoft tensorflow-directml git

https://github.com/microsoft/tensorflow-directml

tensorflo-directmlのインストールと利用

https://hashicco.hatenablog.com/entry/2021/01/18/223516

Scoopの利用方法