Jetson Nanoで自律走行ラジコンを作る(4)

今回はDeep Leaningのモデルを作成し、前回作成したデータセットを使用してモデルを学習させる。

「Jetson Nanoで自律走行ラジコンを作る」シリーズのまとめ記事は以下。

kimitok.hateblo.jp

今回の内容に対応するnotebookは以下(前回の内容も含む)。

github.com

第一回の記事の通り、カメラ画像を入力として前後輪モータの制御値を出力するモデルを考える。モーターの出力は連続値なので回帰問題を解くというイメージになる。

前回記事で作成したデータセットはデータ分割をしていなかったので、アノテーションデータを全部読み込みシャッフルしてから訓練データ・検証データ・テストデータに分割しそれぞれtf.data.Datasetを作成した。今回調べた限りでは、tf.data.Datasetを作成してからデータ分割することは簡単にはできなそうだった。
訓練データ・検証データ・テストデータのサンプル数が8:1:1になるようにデータを分割した。

import random
random.seed(20210616)
random.shuffle(lines)
lines_split = {}
lines_split['train'] = lines[:int(len(lines)*0.8)]
lines_split['val'] = lines[int(len(lines)*0.8):int(len(lines)*0.9)]
lines_split['test'] = lines[int(len(lines)*0.9):]

for type, data in lines_split.items():
  print('{} num: {}'.format(type, len(data)))

train num: 5175
val num: 647
test num: 647

さて、いよいよモデルを設計してみる。
といっても、完全新規でモデルを作成するのは学習が面倒なので、画像から特徴を抽出するバックボーンの部分は画像分類でよく知られているネットワークを持ってきて、その後段(ヘッド)に画像の特徴からモータの制御値を決定する簡単なネットワークを接続する。いわゆる転移学習をする形になる。

Jetson Nano上での処理時間や検出精度などを鑑みて、バックボーンはEfficientNet B0を使用し、ヘッドは640ユニットの全結合層(Denseレイヤ)を2段重ねることにした。
この辺りは後々試行錯誤しそうな気がする。

import tensorflow_hub as hub

num_classes = 2
model = tf.keras.Sequential([
    hub.KerasLayer("https://tfhub.dev/tensorflow/efficientnet/b0/feature-vector/1", trainable=False),
    tf.keras.layers.Dense(640),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.Dense(640),
    tf.keras.layers.BatchNormalization(),
    tf.keras.layers.Activation('relu'),
    tf.keras.layers.Dense(num_classes, activation='linear')
])
model.build([None, 224, 224, 3])  # Batch input shape.
model.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
keras_layer (KerasLayer)     (None, 1280)              4049564   
_________________________________________________________________
dense (Dense)                (None, 640)               819840    
_________________________________________________________________
batch_normalization (BatchNo (None, 640)               2560      
_________________________________________________________________
activation (Activation)      (None, 640)               0         
_________________________________________________________________
dense_1 (Dense)              (None, 640)               410240    
_________________________________________________________________
batch_normalization_1 (Batch (None, 640)               2560      
_________________________________________________________________
activation_1 (Activation)    (None, 640)               0         
_________________________________________________________________
dense_2 (Dense)              (None, 2)                 1282      
=================================================================
Total params: 5,286,046
Trainable params: 1,233,922
Non-trainable params: 4,052,124
_________________________________________________________________

TensorFlow Hubには有名な画像分類器の特徴抽出の部分が単独で公開されているので、画像分類器を持ってきて後段を削除したりしなくてもよい。便利。
今回は回帰問題なので、最後段の活性化関数は分類とは異なりsoftmaxではなく恒等関数(linear)を使っている。

モデルの設定は以下。本当に基本的な設定しかしていない。

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-2),
    loss=tf.keras.losses.mean_squared_error,
    metrics=tf.keras.metrics.mean_absolute_error)

以上でモデルは作れたので学習する。

model.fit(dataset['train'],
          epochs=500,
          steps_per_epoch=len(image_paths['train'])/BATCH_SIZE,
          class_weight={0:3., 1:1.},
          callbacks=[tf.keras.callbacks.TensorBoard(),
                     tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=50, restore_best_weights=True, verbose=1),
                     tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, min_lr=1e-5, patience=20, verbose=1)],
          validation_data=dataset['val'])

自律走行にあたって駆動モーターの制御値よりステアリングモーターの制御値が正確であることが重要ではないかと思ったので、class_weightにより重みに変化をつけている。 あとは検証データの損失を使用して学習率の調整や早期打ち切りの設定をしている。

model.fit()を実行すると粛々とモデルのパラメータが調整される。

Epoch 1/500
161/161 [==============================] - 35s 62ms/step - loss: 0.9651 - mean_absolute_error: 0.4190 - val_loss: 0.1880 - val_mean_absolute_error: 0.3655
Epoch 2/500
161/161 [==============================] - 7s 45ms/step - loss: 0.1396 - mean_absolute_error: 0.2487 - val_loss: 0.0957 - val_mean_absolute_error: 0.2450
Epoch 3/500
161/161 [==============================] - 7s 45ms/step - loss: 0.1280 - mean_absolute_error: 0.2400 - val_loss: 0.0793 - val_mean_absolute_error: 0.2220
Epoch 4/500
161/161 [==============================] - 7s 45ms/step - loss: 0.1028 - mean_absolute_error: 0.2145 - val_loss: 0.0647 - val_mean_absolute_error: 0.1932
Epoch 5/500
161/161 [==============================] - 7s 44ms/step - loss: 0.1007 - mean_absolute_error: 0.2133 - val_loss: 0.0550 - val_mean_absolute_error: 0.1743
...
Epoch 210/500
161/161 [==============================] - 7s 46ms/step - loss: 0.0019 - mean_absolute_error: 0.0303 - val_loss: 0.0302 - val_mean_absolute_error: 0.1067
Epoch 211/500
161/161 [==============================] - 7s 46ms/step - loss: 0.0020 - mean_absolute_error: 0.0310 - val_loss: 0.0302 - val_mean_absolute_error: 0.1068
Epoch 212/500
161/161 [==============================] - 7s 46ms/step - loss: 0.0019 - mean_absolute_error: 0.0302 - val_loss: 0.0301 - val_mean_absolute_error: 0.1065
Epoch 213/500
161/161 [==============================] - 7s 46ms/step - loss: 0.0021 - mean_absolute_error: 0.0309 - val_loss: 0.0303 - val_mean_absolute_error: 0.1063
Epoch 214/500
161/161 [==============================] - 7s 46ms/step - loss: 0.0019 - mean_absolute_error: 0.0300 - val_loss: 0.0303 - val_mean_absolute_error: 0.1065
Epoch 215/500
161/161 [==============================] - 8s 46ms/step - loss: 0.0020 - mean_absolute_error: 0.0306 - val_loss: 0.0302 - val_mean_absolute_error: 0.1060
Restoring model weights from the end of the best epoch.
Epoch 00215: early stopping
<tensorflow.python.keras.callbacks.History at 0x7fb4f9b181d0>

TensorBoardでメトリクスを見てみた感じではまだ精度の改良の余地はありそうに見える。

%load_ext tensorboard
%tensorboard --logdir logs

このモデルを使用して、テストデータに対して推論を実行してみる。

fig = plt.figure(figsize=(24.0, 14.0))
for test_data in dataset['test'].take(1):
  image, label = test_data
  predict_label = model.predict(image)
  for i in range(BATCH_SIZE):
    title = 'actual: {:.2f} {:.2f}\npredict: {:.2f} {:.2f}'.format(label[i][0], label[i][1], predict_label[i][0], predict_label[i][1])

    subplot = fig.add_subplot(4, 8, i+1)
    subplot.imshow(image[i])
    subplot.grid(False)
    subplot.set_title(title)

acutualというのがテストデータの制御値で、predictというのがモデルの推論値である。 ちなみに前回同様画像にはモザイクをかけている。

ぱっと結果をみる限り、それほどぴったり当てているわけではないが、でたらめな値を出力しているわけでもなさそうだ。
ただ、よく見てみるとデータに偏りがある*1ためか、一部データに対する出力の精度が低いように見える。

とはいえ、現状でもそれなりに走りそうな感じがするので、エラー分析は別途するとしてまずはできたモデルをJetson Nanoに組み込んでみようと思う。
推論エンジンとしてTensorRTを使用する想定でモデルをONNXに変換する。
変換にはtf2onnxを使用した。

model.save('models/efficentnet_b0_640_2', overwrite=True)

!pip3 install -U tf2onnx
!python3 -m tf2onnx.convert --saved-model 'models/efficentnet_b0_640_2' --opset 13 --inputs keras_layer_17_input:0[1,224,224,3] --output 'models/efficentnet_b0_640_2.onnx'

念のためONNX Runtimeを使ってONNXのデータに問題がないか軽く動作を確認してみる。

!pip3 install -U onnxruntime

import onnxruntime

sess = onnxruntime.InferenceSession('models/efficentnet_b0_640_2.onnx')

for test_data in dataset['test'].take(1):
  images, labels = test_data
  for i in range(len(images)):
    image = [images[i]]
    label = labels[i]

    pred_onnx = sess.run(None, {'keras_layer_17_input:0' : image})

    print('predict: {:.2f} {:.2f},\n actual: {:.2f} {:.2f}'.format(pred_onnx[0][0][0], pred_onnx[0][0][1], label[0], label[1]))

predict: 0.76 0.81,
 actual: 0.49 1.00
predict: 0.92 1.03,
 actual: 0.87 1.00
predict: 0.09 0.88,
 actual: 0.00 1.00
predict: 0.88 0.55,
 actual: 0.98 0.75
predict: -0.46 0.93,
 actual: -0.50 1.00
...

きちんと動作しているように見える。
今回はここまで。

参考資料

Transfer learning with TensorFlow Hub  |  TensorFlow Core

*1:右旋回で速度レンジが高いデータが多い