PlaidML で Mac 上でも Donkey Car の学習時間を数倍に高速化

この記事は [AI RC Car Adventcalendar 2019](https://qiita.com/advent-calendar/2019/ai-rc-car) の23日目の記事です。

気まぐれで [Donkey Car の GitHub](https://github.com/autorope/donkeycar/)にあるプルリクエストを覗いてみたら、[こんなの](https://github.com/autorope/donkeycar/pull/449)がありました。[PlaidML](https://github.com/plaidml/plaidml)を使って、Mac上での学習時間を削減するものでした。投稿者の Mac では学習時間が 1/3 になったとか。

Mac は かなり前から NVIDIA の GPU の採用を止めていて、Macのユーザのほとんどは、TensorFlow や PyTorch で GPU を使った学習・推論の恩恵に預かることができません。PlaidML は CUDA が無い環境でも GPU を活用した学習・推論ができるよう開発が進められているオープンソースプロジェクトです。PlaidML は Keras のバックエンドにもできるので、ローカル環境での学習を高速化できそうです。

プルリクエストの中身を見たところ、TensorFlow 組み込みの Keras ではなく、生 Keras を使っているので(PlaidML は TensorFloｗ を使わないので、そりゃそうだ)、公式に Donkey Car ソースに取り込まれることはなさそうです。パッチ自体は TensorFlow 組み込みの Keras のインポートから、生 Keras のインポートに書き換えているのがほとんどなので、修正は簡単です。

残念ながらラズパイ用の PlaidML はまだ無かったので、ラズパイでの推論は TensorFlow を使うしかなさそうです。学習と推論で使うバックエンドが違ってくるのでモデルの互換性が心配でしたが、Keras のバージョンが合っていれば大丈夫そうです。

TensorFlow 組み込みの Keras のバーションは以下の方法で確認しました。

```python
$ python 
Python 3.7.3 (default, Mar 27 2019, 16:54:48) 
[Clang 4.0.1 (tags/RELEASE_401/final)] :: Anaconda, Inc. on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> tf.VERSION
'1.14.0'
>>> tf.keras.__version__
'2.2.4-tf'
```

Keras 2.2.4 を使えばよさそうです。

Mac への PlaidML のインストールは [こちら](https://yaritakunai.hatenablog.com/entry/plaidml-test)を参考にしましたが、仮想環境には virtualenv を使わずに conda を使いました。

PlaidML のインストールが終わったら、以下のパッチを Donkey Car の学習環境に適用します。

```patch
diff -ur donkeycar/donkeycar/parts/keras.py donkeycar_plaidml/donkeycar/parts/keras.py
--- donkeycar/donkeycar/parts/keras.py	2019-10-09 23:33:19.000000000 +0900
+++ donkeycar_plaidml/donkeycar/parts/keras.py	2019-10-09 23:39:14.000000000 +0900
@@ -16,15 +16,19 @@
 import numpy as np
 
 import tensorflow as tf
-from tensorflow.python import keras
-from tensorflow.python.keras.layers import Input, Dense
-from tensorflow.python.keras.models import Model, Sequential
-from tensorflow.python.keras.layers import Convolution2D, MaxPooling2D, Reshape, BatchNormalization
-from tensorflow.python.keras.layers import Activation, Dropout, Flatten, Cropping2D, Lambda
-from tensorflow.python.keras.layers.merge import concatenate
-from tensorflow.python.keras.layers import LSTM
-from tensorflow.python.keras.layers.wrappers import TimeDistributed as TD
-from tensorflow.python.keras.layers import Conv3D, MaxPooling3D, Cropping3D, Conv2DTranspose
+
+os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"
+import keras
+from keras.layers import Input, Dense
+from keras.models import Model, Sequential
+from keras.layers import Convolution2D, MaxPooling2D, Reshape, BatchNormalization
+from keras.layers import Activation, Dropout, Flatten, Cropping2D, Lambda
+from keras.layers.merge import concatenate
+from keras.layers import LSTM
+from keras.layers.wrappers import TimeDistributed as TD
+from keras.layers import Conv3D, MaxPooling3D, Cropping3D, Conv2DTranspose
+from keras.utils import CustomObjectScope
+from keras.initializers import glorot_uniform
 
 import donkeycar as dk
 
@@ -48,7 +52,9 @@
         self.optimizer = "adam"
  
     def load(self, model_path):
-        self.model = keras.models.load_model(model_path)
+        #self.model = keras.models.load_model(model_path)
+        with CustomObjectScope({'GlorotUniform': glorot_uniform()}):
+            self.model = keras.models.load_model(model_path)
 
     def load_weights(self, model_path, by_name=True):
         self.model.load_weights(model_path, by_name=by_name)
```

`mycar/train.py` にも以下のパッチを適用します。

```patch
@@ -27,13 +27,15 @@
 import pickle
 import datetime
 
-from tensorflow.python import keras
 from docopt import docopt
 import numpy as np
 from PIL import Image
 
+os.environ["KERAS_BACKEND"] = "plaidml.keras.backend"
+import keras
 import donkeycar as dk
 from donkeycar.parts.datastore import Tub
+
 from donkeycar.parts.keras import KerasLinear, KerasIMU,\
      KerasCategorical, KerasBehavioral, Keras3D_CNN,\
      KerasRNN_LSTM, KerasLatent, KerasLocalizer
```

以上で修正は終わりです。

# 使ってみて

私の Mac のスペックは以下のとおりです。

* MacBook Pro (Retina, 15-inch, Mid 2015)
* CPU 2.5 GHz クアッドコアIntel Core i7
* メモリ 16 GB 1600 MHz DDR3
* GPU AMD Radeon R9 M370X 2 GB

2019/12/1 の [NT札幌 2019](http://wiki.nicotech.jp/nico_tech/index.php?NT%E6%9C%AD%E5%B9%8C2019) での AI カー走行会で、実際にこの PlaidML 環境で学習させてみましたが、出来上がったモデルをラズパイに転送して、無事に自立走行ができました。ベンチマークはとっていなかったのですが CPU での学習に比べて２倍程度は高速化されていました。