PythonでのIoU(交差比)とNMS(非極大値抑制)のGPU加速をcupyで実現

1.はじめに
IoU(交差比)とNMS(非極大値抑制)の計算はターゲット検出タスクにおいて不可欠であるが,計算が必要なbounding boxの数が大きいとcpuは耐えられない.例えばFaster RCNNのRPNネットワークを訓練する際には,約20000個のanchor bboxとground−truth bboxとの間のIoUを計算する必要がある.さらに,RPN出力のproposal RoIから最後のRCNNの訓練サンプルを選択する必要があり,この過程でも10000+のbboxに対してNMS処理を行う必要がある.この場合,その配列演算をGPUに入れて実行し,その並列処理能力を利用して加速を実現する必要がある.
この2つのプロセスを深さ学習フレームワークに組み込んでGPU加速を実現できる場合があるが,Tensorflowのようなこれらのフレームワークで提供される配列タイプで実現するのは容易ではない場合が多いため,本論文で提供する方法は主にCuPyライブラリで提供されるAPIによってIoUとNMSのCUDA加速を実現することである.(ちなみに、CuPyはこれまで深い学習フレームワークChainerの下でcudaベースの下位計算エンジンでしたが、今では独立したライブラリとなり、それに似たライブラリとPyCudaがありました.
2.CuPyのElementwiseKernel類
まず、IoUの計算プロセスをコア関数にカプセル化する必要があります.ここでのコア(Kernel)はCUDAプログラミングの概念であり、私はあまり理解していませんが、GPUで演算を実行する関数として簡単に理解できます.CuPyは、ユーザーに複数のカスタムコア関数の方法を提供しています.具体的には、公式の説明ドキュメントを参照してください.ここでは、CuPyのElementwiseKernelクラスを使用しています.名前の通り、このクラスで定義されたコアによって実行されるのはelement-wiseの演算であり、以下は自分で簡単にまとめたいくつかの注意点です.

各elementの演算は1つの「ループボディ」(loop body)に定義されているが、毎回「ループ」の間は並列である.

「ループ」の回数は出力配列の形状によって決定され、出力配列の形状はElementwiseKernelクラス自身が入力配列の形状に基づいて推定(boardcastingをサポート)してもよいし、ユーザ自身が指定してもよい.

ただし、ループで配列インデックスを使用する必要がある場合は、出力配列の形状を自分で指定する必要があります.

配列は入力されたコアの中でflattenによって1次元の配列になるので、インデックスの場合もflattenの後の位置に従ってインデックスします.

入力コアの変数はすべて読み取り専用で、割り当てはできません.

このクラスおよび他のkernelクラスに関する具体的な構文およびいくつかの簡単な例は、上述した公式説明ドキュメントからも入手できますが、ここでは説明しません.
3.IoUの加速を実現
では、CuPyでelement-wise kernelをカスタマイズする方法を知ったら、GPUで走るIoU関数を書くことができます.

import cupy as cp
import numpy as np

bbox_iou = cp.ElementwiseKernel(
'raw T bbox_a, raw T bbox_b, int32 num_b',
'float32 iou',
'''
//int num_b = sizeof(bbox_b) / (sizeof(bbox_b[0])*8);
int i_a = i / num_b; //bbox_a    ？      
int i_b = i % num_b; //bbox_b    ？      

//bbox_a  i_a     
T width_a = bbox_a[4*i_a + 2] - bbox_a[4*i_a + 0];
T height_a = bbox_a[4*i_a + 3] - bbox_a[4*i_a + 1];
T area_a = width_a * height_a;

//bbox_b  i_b     
T width_b = bbox_b[4*i_b + 2] - bbox_b[4*i_b + 0];
T height_b = bbox_b[4*i_b + 3] - bbox_b[4*i_b + 1];
T area_b = width_b * height_b;

//        
T left = max(bbox_a[4*i_a + 0], bbox_b[4*i_b + 0]);
T top = max(bbox_a[4*i_a + 1], bbox_b[4*i_b + 1]);
T right = min(bbox_a[4*i_a + 2], bbox_b[4*i_b + 2]);
T bottom = min(bbox_a[4*i_a + 3], bbox_b[4*i_b + 3]);
T width_inter = max(right - left, static_cast(0));
T height_inter = max(bottom - top, static_cast(0));
T area_inter = width_inter * height_inter;

iou = area_inter / (area_a + area_b - area_inter + 1e-6);
''',
'IoU')

コードに必要なコメントが入っていますが、どうですか.簡単ですよね?関数は、N個の境界ボックスを含む配列bbox_を実現するaとK個の境界枠を含む配列bbox_bのIoU計算では、出力は(N,K)形の2次元配列であり、実際には[n,k]番目の要素はbbox_を表している.aのn番目のボックスとbbox_bのk番目のボックスのIoU値.
次に、CuPyのarrayでいくつかの簡単な境界ボックスを定義して、送ってみましょう.

a = cp.array([[2, 1, 5, 6],[6, 7, 9, 10], [1, 2, 3, 4]], dtype=cp.float32).reshape(3, 1, 4)
b = cp.array([[2, 1, 5, 6],[6, 7.3, 9, 10.1]], dtype=cp.float32).reshape(1, 2, 4)
y = cp.empty((3, 2), dtype=cp.float32)  #          
out = bbox_iou(a, b, 2, y)  #          bbox   ，     bbox_b     ，     y
print(out.shape)
print(out)

出力は次のとおりです.

(3, 2)

[[ 0.99999994  0.        ]
 [ 0.          0.87096739]
 [ 0.11764705  0.        ]]

4.NMSの加速を実現
非極大値抑制(NMS)の本質は隣接領域の範囲内の非極大値を排除することであり,NMSの原理とアルゴリズムフローについては多くの優れた博文が参考になるが,ここでも後述しない.NMSのCUDA加速の実現には失われた論理的思考が必要だが、難しくはない.NMSアルゴリズムのCPUでの演算のボトルネックは,実際には配列中の現在のbboxと残りのbboxのIoU計算過程にあることを知るだけで,コア関数でこの計算を完了するだけでよいので,NMSの加速は本質的にIoU演算の加速である.前節と唯一異なるのは、2つの異なる配列間のIoUの計算が、単一配列中の各bbox間のIoUの計算になることである.
出力の配列がN個の境界ボックスを含むと仮定すると,我々の出力目標は(N, N)形の配列を出力することである.しかし、i番目のbboxとj番目のbboxのIoUはj番目のbboxとi番目のbboxのIoUと同じであることを知っています.そのため、出力される配列は対称行列(共分散行列に似ていますね)であり、また、あるbboxと自身がIoUを計算するのは意味がないので、この行列の下の三角部分は実際には必要ありません.ゼロにすることができます.この行列を得ると,IoU値が閾値より大きいものを1,閾値より小さいものを0とすることができる.その後、この2値行列に基づいてこれらの境界ボックスの残留を決定することができますが、まずこの行列を出力することを目標にkernel関数を書きましょう.

import cupy as cp
import numpy as np

nms_gpu = cp.ElementwiseKernel(
'raw T bbox, int32 num_box, T threshold',
#'float32 iou',
'uint8 remove_cdd',
'''
int i_x = i % num_box;
int i_y = i / num_box;

// Sinve IoU(a, b) == IoU(b, a) and IoU(a, a) is meaningless, we don't need to 
// calculate the lower triangular part and the diagnal part.
if (i_y < i_x)
{
    T width_x = bbox[4*i_x + 2] - bbox[4*i_x + 0];
    T height_x = bbox[4*i_x + 3] - bbox[4*i_x + 1];
    T area_x = width_x * height_x;
    
    T width_y = bbox[4*i_y + 2] - bbox[4*i_y + 0];
    T height_y = bbox[4*i_y + 3] - bbox[4*i_y + 1];
    T area_y = width_y * height_y;
    
    T left = max(bbox[4*i_x + 0], bbox[4*i_y + 0]);
    T top = max(bbox[4*i_x + 1], bbox[4*i_y + 1]);
    T right = min(bbox[4*i_x + 2], bbox[4*i_y + 2]);
    T bottom = min(bbox[4*i_x + 3], bbox[4*i_y + 3]);
    T width_inter = max(right - left, static_cast(0));
    T height_inter = max(bottom - top, static_cast(0));
    T area_inter = width_inter * height_inter;
    
    float iou = area_inter / (area_x + area_y - area_inter + 1e-6); 
    remove_cdd = (iou > threshold); //    1                       
}

''',
'nms')

コードの内容は、計算を繰り返すIoUと対角線上のIoUを避けるためにif文が複数あるだけで、実際には前節と大同小異であることがわかります.いくつかのbboxを任意に定義し、関数を入力してみましょう.

x = cp.array([[2, 1, 5, 6],[2.6, 1.1, 5, 6], [1, 2, 3, 4], [2.9, 1.1, 5,6], [3, 1.5 ,5.2, 6]], dtype=cp.float32)
y = cp.zeros((num_bbox, num_bbox), dtype=cp.uint8) 
rmv = nms_gpu(x, num_bbox, 0.6, y)
print(rmv)

出力は次のとおりです.

[[0 1 0 1 0]
 [0 0 0 1 1]
 [0 0 0 0 0]
 [0 0 0 0 1]
 [0 0 0 0 0]]

出力マトリクスrmvの[i,j]位置の数値は、i番目のbboxおよびj番目のbboxのIoUがしきい値より大きいかどうかを反映していることがわかる.
次に,最終的にNMSによりbbox長と一致するマスク配列を得ることを望んでおり,その元素が対応する位置をマークしたbboxが除去されるべきかどうかを明らかにする必要がある.次に実現の過程を説明します.皆さんも直接コードを見てもいいし、分かりにくいわけではありません.出力の目標を達成するために、まず長さNの初期値0の1次元配列maskを定義し、maskの値に基づいてrmvを上から下まで逐行判断し(入力したbbox配列は実際には分類されたprobabilityに基づいて大きくから小さく並べ替えられているので、マトリクス処理の優先順位も自然に上から下)、さらにmaskを更新したり演算したりすることができる.例えばrmvマトリクスのi行目に進むと、現在のmaskのi番目の数が0である場合、i番目のbboxが保持される必要があることを示し、マトリクスの行の値はmaskと演算してmaskを更新する.現在のmaskのi番目の数が1である場合、i番目のbboxが削除されることを決定したことを示すと、この行はスキップされ、maskは更新されない.このように,完全なマトリクスを遍歴すると,最終的に必要なマスクが得られる.このコードこそが真の「非極大値抑制」の過程と言える.
以下にコードを示しますが、上のコア関数とは異なり、このコードは実はcpu上で実行されているので、CuPyの配列で書いてみましたが、この順序で実行される論理演算がCPUに適していることが証明されました.もちろんCythonで書くともっと速くなるはずですが、私はCythonができません.興味のある人は自分でやってみてください.

#      
rmv_host = rmv.get() #  cupy array     numpy array，       device   host 
mask = np.zeros(num_bbox, dtype=np.uint8)
for i in range(0, num_bbox-1):
    if not mask[i]:
        mask |= rmv_host[i] # mask rmv_host  i     
print(mask)
idx = np.where(mask==0)[0] #         bbox   
print(idx)

出力は次のとおりです.

[0 1 0 1 0]
[0 2 4]

費用を払う!1行目の出力は私たちの2値maskで、2行目は最後に残されたbboxのインデックスです.
このコードは実際にFaster RCNNの作者のソースコードを参考にしていますが、原作者の考え方は本当に牛が多く、64個のbboxごとにGPUを割り当てるblockの中で、64ビットULL型のバイナリ数で二値maskを表し、ビットやmaskを更新する...実に巧みです.Faster RCNNのNMSソース部分の分析については、このブログを見てみると、目標検出におけるNMSのGPU実装(Faster R-CNNのnms_kernel.cuファイルから).
最後に、GPUによってNMSが加速する効果を検証し、numpy配列に基づいてcpuで実行されるnms関数を定義します(実際には、IoUと比較しきい値を計算する部分のみを定義します):

def nms_cpu(bbox, num_box, threshold):
    remove = np.zeros((num_box, num_box), dtype=np.uint8)
    for i in range(num_box):
        for j in range(num_box):  
            if i < j:
                width_x = bbox[j, 2] - bbox[j, 0]
                height_x = bbox[j, 3] - bbox[j, 1]
                area_x = width_x * height_x
                
                width_y = bbox[i, 2] - bbox[i, 0]
                height_y = bbox[i, 3] - bbox[i, 1]
                area_y = width_y * height_y
                
                left = max(bbox[j, 0], bbox[i, 0])
                top = max(bbox[j, 1], bbox[i, 1])
                right = min(bbox[j, 2], bbox[i, 2])
                bottom = min(bbox[j, 3], bbox[i, 3])
                width_inter = max(right - left, 0.)
                height_inter = max(bottom - top, 0.)
                area_inter = width_inter * height_inter
                
                iou = area_inter / (area_x + area_y - area_inter + 1e-6) 
                remove[i, j] = (iou > threshold)
    return remove

次に、Numpy arrayとCuPy arrayを使用して、5000個のbboxを含む配列をそれぞれ定義し(配列要素をすべて1に設定しやすいように)、CPUの関数とGPUの関数にそれぞれ送り、効果を見てみましょう.

#      ，      mask       ，       iou     
#     ，   win10  ，gpu 1080TI，cpu i7-8700
import time

num_bbox = 5000

s = time.time()
a = np.ones((num_bbox, 4), dtype=np.float32) 
rmv_2 = nms_cpu(a, num_bbox, 0.5)
e = time.time()
print('   ：', e - s, 's')

s = time.time()
b = cp.ones((num_bbox, 4), dtype=cp.float32) 
z = cp.zeros((num_bbox, num_bbox), dtype=cp.uint8) 
rmv_1 = nms(b, num_bbox, 0.5, z)
e = time.time()
print('   ：', e - s, 's')

出力結果:

   ： 96.38717579841614 s
   ： 0.006020069122314453 s

寝槽、直接離陸してありますか!?
5.後記
最近他の博文で多くのことを学んで、それから多くの博主は私が残した問題と評論に対してすべて辛抱強い解答を行って、とても彼らに感謝して、そこで自分のいくつかの進展と経験を分かち合いたいと思って、いくつかの技術の含有量がないものですが、しかしみんなにも多少助けてほしいです.また、私自身も初心者です.Pythonのプログラミングにおいても、深い学習においても、文章に漏れや間違いがあれば、ご指摘を歓迎します.