前回、プロファイリングを行った結果として、nmsという関数の処理に時間がかかっていることがわかったというお話をしました。

前回もご説明しましたが、NMSは簡単に言えば複数ある候補から一番いいものを選択するためのアルゴリズムです。そういうわけでNMS自体はディープラーニングとも機械学習とも直接は関係ないのですが、推論結果の処理をする際に比較的頻繁に用いられます。ちなみに私が最初にNMSに触れたのは、産業ロボットの画像処理プログラムでエッジ検知アルゴリズム(Canny法)を実装した時でした。エッジの尾根(極大部)以外を抑制するアルゴリズムとしてNMSを使ったのです。

SSDやYoloのような特徴マップにアンカーボックスを配置する方法では、ひとつのオブジェクトに対し候補となるバウンディングボックスが複数作られます。この中からNMSを用いてもっとも有力なバウンディングボックスをひとつだけ採用するのですが、特にアンカーボックスを多く配置した場合には候補のバウンディングボックスが大量に算出されます。結果としてNMSの処理に多大な計算量が消費されることになるのです。

実はNMSは、実装によってこの処理量に大きな差がでるのです。このことはリアルタイム検知を行う場合や、特に組込み機器のような計算リソースが限られているハードウェアに実装する場合には重要なポイントとなります。

物体検知の実装の中にはNMSの部分だけCで実装したり、GPUを用いた実装としているものが多くあります。対してこれまで公開してきた物体検知のNMSはPythonで実装したもので、比較的低速です。

調べてみたところ、Pythonで工夫してNMSを高速に実装した例を見つけました。オリジナルのページには高速化のポイントが以下のように書かれています。

Instead, Dr. Malisiewicz replaced this inner for loop with vectorized code — and this is how we are able to achieve substantially faster speeds when applying non-maximum suppression.

要するにforループの代わりにvector演算を使うことによって高速化を実現したということですが、これはPythonのnumpyライブラリを用いる高速化テクニックの定番です。

ただしこの実装ではIoUを比較対象にしていないためにその部分を書き直し、これまでのコードのnmsと同じ使い方で使えるように調整しました。以下がそのコードになります。

def fast_nms(boxes, scores, nms_thresh=0.5, top_k=200):
    boxes = boxes.cpu().numpy()
    scores = scores.cpu().numpy()
    keep = []
    if len(boxes) == 0:
        return keep
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = (x2-x1)*(y2-y1)
    idx = np.argsort(scores, axis=0)   # sort in ascending order
    idx = idx[-top_k:]  # indices of the top-k largest vals

    count = 0
    while len(idx) > 0:
        last = len(idx)-1
        i = idx[last]  # index of current largest val
        keep.append(i)
        count  = 1
        xx1 = np.maximum(x1[i], x1[idx[:last]])
        yy1 = np.maximum(y1[i], y1[idx[:last]])
        xx2 = np.minimum(x2[i], x2[idx[:last]])
        yy2 = np.minimum(y2[i], y2[idx[:last]])

        w = np.maximum(0, xx2-xx1)
        h = np.maximum(0, yy2-yy1)

        inter = w*h
        iou = inter / (area[idx[:last]] area[i]-inter)
        idx = np.delete(idx, np.concatenate(([last], np.where(iou > nms_thresh)[0])))

    return keep, count

高速化に成功しているか、あらためてプロファイリングを行ってみます。

  $ CUDA_VISIBLE_DEVICES=0 python -m cProfile -s time -o profile.txt tools/demo_mp.py --cfg experiments/coco/efficientnet/efficient_b1_256x192_40_32_32_32_adam_lr1e-3.yaml
  $ echo -e 'sort tottime\nstats' | python3 -m pstats profile.txt  | less
  Welcome to the profile statistics browser.
  profile.txt% profile.txt% Sat Aug 15 13:25:35 2020    profile.txt
  1352533 function calls (1319451 primitive calls) in 18.400 seconds
  Ordered by: internal time
  ncalls  tottime  percall  cumtime  percall filename:lineno(function)
  2994    3.575    0.001    3.575    0.001 {method 'acquire' of '_thread.lock' objects}
  3544    2.341    0.001    2.341    0.001 {built-in method __new__ of type object at 0x563940a10d40}
  1    2.204    2.204    4.309    4.309 /home/eigo/src/temp/posenet-demo-trt_1Q-fast-nms/venv/lib/python3.7/site-
  packages/imutils/video/webcamvideostream.py:6(__init__)
  1    2.105    2.105    2.105    2.105 {method 'read' of 'cv2.VideoCapture' objects}
  115/112    1.312    0.011    1.315    0.012 {built-in method _imp.create_dynamic}
  1    1.001    1.001    1.001    1.001 {built-in method time.sleep}
  206    0.728    0.004    7.328    0.036 tools/pose_util_mp.py:204(predict)
  206    0.690    0.003    0.690    0.003 {method 'cuda' of 'torch._C._TensorBase' objects}
  206    0.597    0.003    0.597    0.003 {waitKey}
  838    0.582    0.001    0.582    0.001 {method 'cpu' of 'torch._C._TensorBase' objects}
  205    0.452    0.002    0.452    0.002 {imshow}
  206    0.322    0.002    0.322    0.002 {method 'argsort' of 'numpy.ndarray' objects}
  266    0.265    0.001    0.265    0.001 {method 'copy' of 'numpy.ndarray' objects}
  206    0.220    0.001    0.796    0.004 /home/eigo/src/temp/posenet-demo-trt_1Q-fast-nms/posenet/object_detector_efnetssd/utils/box_utils.py:252(fast_nms)
  823    0.214    0.000    0.214    0.000 {built-in method cat}
  442    0.133    0.000    0.133    0.000 {method 'astype' of 'numpy.ndarray' objects}
  412    0.131    0.000    0.131    0.000 {method 'to' of 'torch._C._TensorBase' objects}
  974    0.118    0.000    0.118    0.000 {built-in method marshal.loads}
...

percallの値を見ると数倍の速度にはなっているようです。GitHubのコードにも反映しておきましたのでご参考にしてください。