Darknet/YOLOのJetson最適化

Mixed Precision学習結果

ソースコードはこちら！

表題の通り、Darknet/YOLOv3をJetson Nanoに最適化しました。

本最適化により、Jetson Nanoでより少ないメモリ利用で学習・検出ができるようになりました。
すなわち、通常Darknet/YOLOv3に対して、より大きい解像度で学習・検出ができます。
学習は二つのクラス 320x320で11日程(416x416で14.5日)でとても遅いけど、Jetson Nanoしか持っていない方は自分のデータで学習ができます。

Jetson Nanoでしか動作確認をしていないが、TX1/2等のMaxwellアーキテクチャを実装するJetsonでは動作すると思います。
Xavierは一部のコード（AtomicAdd）で正しくアーキテクチャの切替ができるのか分からないです。
一方、通常のNvidia GPUカードでは未対応です（現状、JetsonのUnified Memoryに特化）。

最適化内容は下記二つです：
・JetsonのUnified Memoryを利用 
・FP32/16混合学習(Mixed Precision) (論文、実装、実装、NVidia)

状況：
・テストは最低限しか行っていないです。
・対応層：convolutional, route, shortcut, upsample, yolo, cost.通常YOLOv3しか動作しないです。v4, v2やTinyは動作しないです。

cfg設定ファイルへのMixed Precision追加項目（[net]部分のみ）は以下です。説明は設定項目名の通りです。詳細は後述のアルゴリズムを参考のこと。

追加項目名	デフォルト値	備考
max_loss_scale	16384.0
min_loss_scale	.0009765625
initial_loss_scale	1.0
loss_scale_inc_iterations	1000	burn_in以降に利用

コンパイル
Makefileに二つ項目を追加しました：
* JETMEM: JetsonのUnified Memoryを利用の際、1にします。
* CUDNN_MIX: FP16/Mixed Precisionを利用の際、1にします。
現状、JETMEM=1, CUDNN_MIX=1でしか動作確認していないです。。
また、Jetson用のARCHのコメントも外します。（compute_53を利用）

検出
既存の重み、または独自データで学習した重みが利用できます。
使い方はこちらとこちら。

学習
学習の実験はCOCOデータセットの猫と犬のみで行いました。
YOLOv3用のCOCOデータセットの取得はこちらに従ってから行いましたが、二つのクラスのみを使うために新規設定ファイルを作成しました：

data/coco/coco_cat_dog.names　猫と犬のみのクラス名
cfg/coco_cat_dog.data　学習および評価ファイルリストを更新（下記二つ）。
data/coco/trainval_cat_dog.txt　学習用ファイルリスト
data/coco/val_cat_dog.txt　評価用ファイルリスト
cfg/yolov3_cat_dog*.cfg　二つのクラスのみを利用するためにフィルタ数対応

また、新しいラベルに対応するファイルも必要です。以下を行います。（注意：上記の通りに取得したCOCOデータセットを上書きします。）

  tar xvf coco_cat_dog_labels.tar.gz

学習例は以下です：
./darknet detector train cfg/coco_cat_dog.data cfg/yolov3_cat_dog.cfg weights/darknet53.conv.74

最適化実装の詳細：

Unified Memory
Darknetは、CPUとGPUに同じデータを保存します。しかし、JetsonのDRAMは両CPU/GPUで共有されるため、同じデータを２回確保する必要はないです。
よって、共有することにより、より少ないメモリで動作が可能になり、より大きい解像度で検出が可能です。

FP32/16混合学習(Mixed Precision)
Mixed Precisionは学習の際にNVidia GPUのTensorCore（Turingアーキテクチャ以降）を利用して高速化を達成するために開発されました。FP32に対し、FP16での畳込み演算のみが８倍速くなります。FP16を用いることにより、利用するメモリ量も減ります。
（重み更新用の配列等はFP32/16の両方を保持する必要があるため増えますが、全体は減ります）
Jetson NanoはTensorCoreを実現しないため学習は遅くなりますが（後述）、より少ないメモリ（より大きい解像度）の学習が可能になります。

 演算部の実装はこちらのTensorFlow部分を元に行いました。

論文のように、FP32⇔16変換を行う際、FP16はより少ない数値幅のためloss scaleで乗算/除算を行いますがこれの調整が必要になります。
更新案はこちらやこちらにあり、基本は重みを更新するweight_updatesがNaNまたは無限になった場合にloss scaleを半分にし、
NaN/無限が 1000または2000回生じなかった場合に倍にしています。
ここでは倍にする手順は上記と同じにしていますが、YOLOv3の特徴によりburn_in（YOLOv3のcfgファイルでは1000に設定されています）回までは最大の
weight_updatesにより更新し、burn_in後は従来の通りにしています。理由や詳細は次の解析の通りです。

上記図は各回毎の、全て畳込層のweight_updatesの最大絶対値グラフの形よびロスのFP32例です。
パラメータ調整や、二回目の学習だと別結果になりますが、雰囲気は同じです。
最初はすぐ大きくなります。この例だと、最大が165,055です。この時点でFP16の最大値（65,504）を超えています。
よって、この時点でloss scaleの適切な値は0.25です。その後、徐々に下がり、数百回後に2^10程下がります。
これにより、1000回に一回倍にするとFP32⇔16では精度を失う可能性があるため、より早く更新するために今回および今回までのweight_updatesの最大絶対値に基づいて
burn_inまで更新するようにしました。

また、下記が層毎のweight_updatesです。（最大絶対値が大きい場合）
ここでは５００回目程度まで、一番大きい絶対値は層81、93、105です。これらはlossが計算されるyolo_layerの直前の畳み込み層です。
５００回目以降になると、最も大きい絶対値は最初の方の層になります。

大きさ変更による水増し
Darknet/YOLOv3のデータ水増しとして、色調整や、大きさ変更（Resize）、横反転、切り出しがあります。
Resizeは１０回毎に、ランダムに1/1.4～1.4倍に、width/heightが32の倍数で調整されまる。
320x320の場合、256x256から480x480に変更されますが、最初はメモリを最大で初期化するため、必ず最大の480x480になります。

メモリ使用量評価方法

Tegraでは下記方法があります：
* tegrastats コマンド（メモリ使用量を指定周期で指定ファイルに出力可能）
* jtop コマンド(tegrastatsと同じ情報であるが、ファイルに出力できない)
* cudaMemGetInfo (c ファイルに埋め込み)
* /usr/bin/time -v (Maximum resident set sizeを参考)

ここではファイルに指定周期で出力可能なtegrastatsを使うことにしました。
なお、cudaMemGetInfo, /usr/bin/time, tegrastatsは別数値を出力します。

Nano上での実験結果
以下が学習中のメモリ使用量です。
メモリ使用量はtegrastatsによる RAM + SWAPであり、他のプロセスのメモリ使用量も含まれています。
上記のように、Darknet/YOLOv3の最大メモリ使用量はResizeが最大の場合であり、最初の１０回はこのResizeに設定されます。
よって、最初の１０回のメモリ使用量が最大になること考えられるけど、学習中、他のプロセスのため全体のメモリ使用量はさらに増えます。
このため、Darknetは下記二つの理由により強制終了（Killed）される可能性があります：
1) 最初のメモリ確保の場合
2) 学習中、他のプロセスによるメモリ増加のため

例えば以下が320x320のメモリ使用量です。
FP32の方が最初に完了します。最初の１０回のメモリ使用量(RAM + SWAP)は 3800MB(FP32)および3200MB(Mixed Precision)です。
しかし、学習中、メモリ使用量はさらに増えますが、解析が必要です。

なお、上記のNanoの設定として、SWAPがデフォルトの2GBと、ターミナルモードです。

以下が学習一覧です。

設定	元Darknet/YOLOv3	本最適化
yolov3 320x320	OK	OK
yolov3 416x416	100回程度で強制終了	OK
yolov3 544x544	初期に強制終了	300回程度で強制終了
yolov3 tiny	解像度依存	未実装
yolov2	解像度依存	未実装

実行時間解析

320x320設定で、Resize 480x480の一回の学習にかかる時間は以下です。

  FP32 : 186 秒

  Mix : 272 秒 

以下の情報は次のコマンドのログからまとめています。

  nvprof --print-gpu-trace --log-file log.txt ./darknet detector train ...

　上記で、学習を３回繰り返し、二回目の結果を分析しました。

　FP32の場合、GPUの合計時間は、上記186秒とほぼ同じです、

　しかし、Mixed Precisionでは、GPU時間が232秒であり、上記272秒より40秒少ないです。

　一方、私の解析が正しければ、下記のようにMixとFP32で最も大きい差分は畳み込み演算です（Mixが47秒遅い）。よって、畳み込み演算が大きな差分でありますが、さらに40秒の解析が必要です。

mAP@0.50%結果（猫、犬のみ）

w x h	モード	最良mAP
320x320	2クラス (FP32)	66.4%
320x320	2 クラス (Mixed Precision)	68.0%
416x416	2 クラス (Mixed Precision)	73.3%

ここでは5000回学習し、最も良いmAPの重みになります（100回毎に重みを保存）。
mAPは同じ設定でも、学習毎にかわります。
よって、上記結果はMixed Precision方式がFP32より良いと限りません。