Unity ML-AgentsをAMD製GPU + ROCm + TensorFlowで使う

Unity ML-AgentsはTensorFlowで実装されているんですが、RadeonでTensorFlowを使うのは色々と面倒です。でも、RadeonでML-Agentsを使いたかったのでチャレンジしてみました。

結論としては、RadeonでML-Agentsを動作させることができて、radeontopの表示でばっちりGPU使用してくれました。が、なぜか学習速度がCPUより遅くて実用性がないです。ML-Agentsには画像を入力して学習させたのですが、解像度が低くても高くてもCPUより処理が遅くなりました。

NVIDIAのGPU インスタンスも借りてTeslaでも試してみたのですが、こちらもCPUより遅かったです。ML-AgentsはGPUと相性が悪いのでしょうか？GPUにデータ流し込むのに時間がかかっているのでしょうかね。

ML-Agentsをやる場合はGPUよりCPUに投資する方がいいかもしれません。

以前の問題点

ML-AgentsがTensorFlow 2.0に対応したので今は普通にROCmで動作させられると思います。（私はまだ未確認）

~~ML-AgentsはTensorFlow 1.7.1という結構古いバージョンを必要としますが、ROCmが対応するTensorFlowにはそんな古いバージョンがありません。~~

仕方ないので今回はTensorFlow 1.15.0rc2をインストールして無理やり動かしました。そのうちML-AgentsもTensorFlow 2.0に対応すると思うので、無理せず待つのがいいでしょうね。

ROCmは動作するハードウェアが非常に限られています。確認しましょう。 github.com

Ubuntu 18.04.3 LTSにしました。他のディストリビューションだとインストール手順が書かれてないので難しそうです。

0.10.1を使用しました。

個人的にml-agentsはPythonの仮想環境で使用したいので、virtualenvをインストールしました。

sudo apt install virtualenv

TensorFlow最新版を使用するのであれば公式サイトの手順どおりにコマンドを打てばOKですが、TensorFlow 1.15.0rc2に必要なライブラリであるrcclがインストールされなかったので別途インストールしました。

ROCmはRadeon公式ドライバがインストールされていると動作に支障がでる？らしいです。私は画面表示はIntel内臓グラフィックスを使用しました。

ML-Agentsを普通にインストールすると、ROCmのTensorFlowが上書きされてしまいます。ML-Agentsインストール後に、TensorFlowをアンインストールして、 tensorflow-rocmをインストールしましょう。

pip install mlagents==0.10.1
pip uninstall tensorflow
pip install tensorflow-rocm==1.15.0rc2

TensorFlowとML-Agentsが両方GPUで動いているか確認しましょう。TensorFlowのバージョンによっては動かないです。

今回はML-AgentsをAMD製GPUで動かせることはわかりました。ML-AgentsがTensorFlow 2.0に対応してから再度試したいと思います。